lakehouse_engine.algorithms.reconciliator

  1"""Module containing the Reconciliator class."""
  2
  3from enum import Enum
  4from typing import List
  5
  6import pyspark.sql.functions as spark_fns
  7from pyspark.sql import DataFrame
  8from pyspark.sql.functions import abs, coalesce, col, lit, when  # noqa: A004
  9from pyspark.sql.types import FloatType
 10
 11from lakehouse_engine.algorithms.exceptions import ReconciliationFailedException
 12from lakehouse_engine.core.definitions import InputSpec, ReconciliatorSpec
 13from lakehouse_engine.core.exec_env import ExecEnv
 14from lakehouse_engine.core.executable import Executable
 15from lakehouse_engine.io.reader_factory import ReaderFactory
 16from lakehouse_engine.transformers.optimizers import Optimizers
 17from lakehouse_engine.utils.logging_handler import LoggingHandler
 18
 19
 20class ReconciliationType(Enum):
 21    """Type of Reconciliation."""
 22
 23    PCT = "percentage"
 24    ABS = "absolute"
 25
 26
 27class ReconciliationTransformers(Enum):
 28    """Transformers Available for the Reconciliation Algorithm."""
 29
 30    AVAILABLE_TRANSFORMERS: dict = {
 31        "cache": Optimizers.cache,
 32        "persist": Optimizers.persist,
 33    }
 34
 35
 36class Reconciliator(Executable):
 37    """Class to define the behavior of an algorithm that checks if data reconciles.
 38
 39    Checking if data reconciles, using this algorithm, is a matter of reading the
 40    'truth' data and the 'current' data. You can use any input specification compatible
 41    with the lakehouse engine to read 'truth' or 'current' data. On top of that, you
 42    can pass a 'truth_preprocess_query' and a 'current_preprocess_query' so you can
 43    preprocess the data before it goes into the actual reconciliation process.
 44    Moreover, you can use the 'truth_preprocess_query_args' and
 45    'current_preprocess_query_args' to pass additional arguments to be used to apply
 46    additional operations on top of the dataframe, resulting from the previous steps.
 47    With these arguments you can apply additional operations like caching or persisting
 48    the Dataframe. The way to pass the additional arguments for the operations is
 49    similar to the TransformSpec, but only a few operations are allowed. Those are
 50    defined in ReconciliationTransformers.AVAILABLE_TRANSFORMERS.
 51
 52    The reconciliation process is focused on joining 'truth' with 'current' by all
 53    provided columns except the ones passed as 'metrics'. After that it calculates the
 54    differences in the metrics attributes (either percentage or absolute difference).
 55    Finally, it aggregates the differences, using the supplied aggregation function
 56    (e.g., sum, avg, min, max, etc).
 57
 58    All of these configurations are passed via the ACON to instantiate a
 59    ReconciliatorSpec object.
 60
 61    .. note::
 62        It is crucial that both the current and truth datasets have exactly the same
 63        structure.
 64    .. note::
 65        You should not use 0 as yellow or red threshold, as the algorithm will verify
 66        if the difference between the truth and current values is bigger
 67        or equal than those thresholds.
 68    .. note::
 69        The reconciliation does not produce any negative values or percentages, as we
 70        use the absolute value of the differences. This means that the recon result
 71        will not indicate if it was the current values that were bigger or smaller
 72        than the truth values, or vice versa.
 73    """
 74
 75    _logger = LoggingHandler(__name__).get_logger()
 76
 77    def __init__(self, acon: dict):
 78        """Construct Algorithm instances.
 79
 80        Args:
 81            acon: algorithm configuration.
 82        """
 83        self.spec: ReconciliatorSpec = ReconciliatorSpec(
 84            metrics=acon["metrics"],
 85            truth_input_spec=InputSpec(**acon["truth_input_spec"]),
 86            current_input_spec=InputSpec(**acon["current_input_spec"]),
 87            truth_preprocess_query=acon.get("truth_preprocess_query", None),
 88            truth_preprocess_query_args=acon.get("truth_preprocess_query_args", None),
 89            current_preprocess_query=acon.get("current_preprocess_query", None),
 90            current_preprocess_query_args=acon.get(
 91                "current_preprocess_query_args", None
 92            ),
 93            ignore_empty_df=acon.get("ignore_empty_df", False),
 94        )
 95
 96    def get_source_of_truth(self) -> DataFrame:
 97        """Get the source of truth (expected result) for the reconciliation process.
 98
 99        Returns:
100            DataFrame containing the source of truth.
101        """
102        truth_df = ReaderFactory.get_data(self.spec.truth_input_spec)
103        if self.spec.truth_preprocess_query:
104            truth_df.createOrReplaceTempView("truth")
105            truth_df = ExecEnv.SESSION.sql(self.spec.truth_preprocess_query)
106
107        return truth_df
108
109    def get_current_results(self) -> DataFrame:
110        """Get the current results from the table that we are checking if it reconciles.
111
112        Returns:
113            DataFrame containing the current results.
114        """
115        current_df = ReaderFactory.get_data(self.spec.current_input_spec)
116        if self.spec.current_preprocess_query:
117            current_df.createOrReplaceTempView("current")
118            current_df = ExecEnv.SESSION.sql(self.spec.current_preprocess_query)
119
120        return current_df
121
122    def execute(self) -> None:
123        """Reconcile the current results against the truth dataset."""
124        truth_df = self.get_source_of_truth()
125        self._apply_preprocess_query_args(
126            truth_df, self.spec.truth_preprocess_query_args
127        )
128        self._logger.info("Source of truth:")
129        truth_df.show(1000, truncate=False)
130
131        current_results_df = self.get_current_results()
132        self._apply_preprocess_query_args(
133            current_results_df, self.spec.current_preprocess_query_args
134        )
135        self._logger.info("Current results:")
136        current_results_df.show(1000, truncate=False)
137
138        status = "green"
139
140        # if ignore_empty_df is true, run empty check on truth_df and current_results_df
141        # if both the dataframes are empty then exit with green
142        if (
143            self.spec.ignore_empty_df
144            and truth_df.isEmpty()
145            and current_results_df.isEmpty()
146        ):
147            self._logger.info(
148                f"ignore_empty_df is {self.spec.ignore_empty_df}, "
149                f"truth_df and current_results_df are empty, "
150                f"hence ignoring reconciliation"
151            )
152            self._logger.info("The Reconciliation process has succeeded.")
153            return
154
155        recon_results = self._get_recon_results(
156            truth_df, current_results_df, self.spec.metrics
157        )
158        self._logger.info(f"Reconciliation result: {recon_results}")
159
160        for m in self.spec.metrics:
161            metric_name = f"{m['metric']}_{m['type']}_diff_{m['aggregation']}"
162            if m["yellow"] <= recon_results[metric_name] < m["red"]:
163                if status == "green":
164                    # only switch to yellow if it was green before, otherwise we want
165                    # to preserve 'red' as the final status.
166                    status = "yellow"
167            elif m["red"] <= recon_results[metric_name]:
168                status = "red"
169
170        if status != "green":
171            raise ReconciliationFailedException(
172                f"The Reconciliation process has failed with status: {status}."
173            )
174        else:
175            self._logger.info("The Reconciliation process has succeeded.")
176
177    @staticmethod
178    def _apply_preprocess_query_args(
179        df: DataFrame, preprocess_query_args: List[dict]
180    ) -> DataFrame:
181        """Apply transformers on top of the preprocessed query.
182
183        Args:
184            df: dataframe being transformed.
185            preprocess_query_args: dict having the functions/transformations to
186                apply and respective arguments.
187
188        Returns: the transformed Dataframe.
189        """
190        transformed_df = df
191
192        if preprocess_query_args is None:
193            transformed_df = df.transform(Optimizers.cache())
194        elif len(preprocess_query_args) > 0:
195            for transformation in preprocess_query_args:
196                rec_func = ReconciliationTransformers.AVAILABLE_TRANSFORMERS.value[
197                    transformation["function"]
198                ](**transformation.get("args", {}))
199
200                transformed_df = df.transform(rec_func)
201        else:
202            transformed_df = df
203
204        return transformed_df
205
206    def _get_recon_results(
207        self, truth_df: DataFrame, current_results_df: DataFrame, metrics: List[dict]
208    ) -> dict:
209        """Get the reconciliation results by comparing truth_df with current_results_df.
210
211        Args:
212            truth_df: dataframe with the truth data to reconcile against. It is
213                typically an aggregated dataset to use as baseline and then we match the
214                current_results_df (Aggregated at the same level) against this truth.
215            current_results_df: dataframe with the current results of the dataset we
216                are trying to reconcile.
217            metrics: list of dicts containing metric, aggregation, yellow threshold and
218                red threshold.
219
220        Return:
221            dictionary with the results (difference between truth and current results)
222        """
223        if len(truth_df.head(1)) == 0 or len(current_results_df.head(1)) == 0:
224            raise ReconciliationFailedException(
225                "The reconciliation has failed because either the truth dataset or the "
226                "current results dataset was empty."
227            )
228
229        # truth and current are joined on all columns except the metrics
230        joined_df = truth_df.alias("truth").join(
231            current_results_df.alias("current"),
232            [
233                truth_df[c] == current_results_df[c]
234                for c in current_results_df.columns
235                if c not in [m["metric"] for m in metrics]
236            ],
237            how="full",
238        )
239
240        for m in metrics:
241            if m["type"] == ReconciliationType.PCT.value:
242                joined_df = joined_df.withColumn(
243                    f"{m['metric']}_{m['type']}_diff",
244                    coalesce(
245                        (
246                            # we need to make sure we don't produce negative values
247                            # because our thresholds only accept > or >= comparisons.
248                            abs(
249                                (
250                                    col(f"current.{m['metric']}")
251                                    - col(f"truth.{m['metric']}")
252                                )
253                                / abs(col(f"truth.{m['metric']}"))
254                            )
255                        ),
256                        # if the formula above produces null, we need to consider where
257                        # it came from: we check below if the values were the same,
258                        # and if so the diff is 0, if not the diff is 1 (e.g., the null
259                        # result might have come from a division by 0).
260                        when(
261                            col(f"current.{m['metric']}").eqNullSafe(
262                                col(f"truth.{m['metric']}")
263                            ),
264                            lit(0),
265                        ).otherwise(lit(1)),
266                    ),
267                )
268            elif m["type"] == ReconciliationType.ABS.value:
269                joined_df = joined_df.withColumn(
270                    f"{m['metric']}_{m['type']}_diff",
271                    abs(
272                        coalesce(col(f"current.{m['metric']}"), lit(0))
273                        - coalesce(col(f"truth.{m['metric']}"), lit(0))
274                    ),
275                )
276            else:
277                raise NotImplementedError(
278                    "The requested reconciliation type is not yet implemented."
279                )
280
281            joined_df = joined_df.withColumn(
282                f"{m['metric']}_{m['type']}_diff",
283                col(f"{m['metric']}_{m['type']}_diff").cast(FloatType()),
284            )
285
286        results_df = joined_df.agg(
287            *[
288                getattr(spark_fns, m["aggregation"])(
289                    f"{m['metric']}_{m['type']}_diff"
290                ).alias(f"{m['metric']}_{m['type']}_diff_{m['aggregation']}")
291                for m in metrics
292            ]
293        )
294
295        return results_df.collect()[0].asDict()