lakehouse_engine.transformers.filters

  1"""Module containing the filters transformers."""
  2
  3from typing import Any, Callable, List, Optional
  4
  5from pyspark.sql import DataFrame
  6from pyspark.sql.functions import col
  7
  8from lakehouse_engine.transformers.watermarker import Watermarker
  9from lakehouse_engine.utils.logging_handler import LoggingHandler
 10
 11
 12class Filters(object):
 13    """Class containing the filters transformers."""
 14
 15    _logger = LoggingHandler(__name__).get_logger()
 16
 17    @classmethod
 18    def incremental_filter(
 19        cls,
 20        input_col: str,
 21        increment_value: Optional[Any] = None,
 22        increment_df: Optional[DataFrame] = None,
 23        increment_col: str = "latest",
 24        greater_or_equal: bool = False,
 25    ) -> Callable:
 26        """Incrementally Filter a certain dataframe given an increment logic.
 27
 28        This logic can either be an increment value or an increment dataframe from
 29        which the get the latest value from. By default, the operator for the
 30        filtering process is greater or equal to cover cases where we receive late
 31        arriving data not cover in a previous load. You can change greater_or_equal
 32        to false to use greater, when you trust the source will never output more data
 33        with the increment after you have load the data (e.g., you will never load
 34        data until the source is still dumping data, which may cause you to get an
 35        incomplete picture of the last arrived data).
 36
 37        Args:
 38            input_col: input column name
 39            increment_value: value to which to filter the data, considering the
 40                provided input_Col.
 41            increment_df: a dataframe to get the increment value from.
 42                you either specify this or the increment_value (this takes precedence).
 43                This is a good approach to get the latest value from a given dataframe
 44                that was read and apply that value as filter here. In this way you can
 45                perform incremental loads based on the last value of a given dataframe
 46                (e.g., table or file based). Can be used together with the
 47                get_max_value transformer to accomplish these incremental based loads.
 48                See our append load feature tests  to see how to provide an acon for
 49                incremental loads, taking advantage of the scenario explained here.
 50            increment_col: name of the column from which to get the increment
 51                value from (when using increment_df approach). This assumes there's
 52                only one row in the increment_df, reason why is a good idea to use
 53                together with the get_max_value transformer. Defaults to "latest"
 54                because that's the default output column name provided by the
 55                get_max_value transformer.
 56            greater_or_equal: if filtering should be done by also including the
 57                increment value or not (useful for scenarios where you are performing
 58                increment loads but still want to include data considering the increment
 59                value, and not only values greater than that increment... examples may
 60                include scenarios where you already loaded data including those values,
 61                but the source produced more data containing those values).
 62                Defaults to false.
 63
 64        Returns:
 65            A function to be called in .transform() spark function.
 66        """
 67
 68        def inner(df: DataFrame) -> DataFrame:
 69            if increment_df:
 70                if greater_or_equal:
 71                    return df.filter(  # type: ignore
 72                        col(input_col) >= increment_df.collect()[0][increment_col]
 73                    )
 74                else:
 75                    return df.filter(  # type: ignore
 76                        col(input_col) > increment_df.collect()[0][increment_col]
 77                    )
 78            else:
 79                if greater_or_equal:
 80                    return df.filter(col(input_col) >= increment_value)  # type: ignore
 81                else:
 82                    return df.filter(col(input_col) > increment_value)  # type: ignore
 83
 84        return inner
 85
 86    @staticmethod
 87    def expression_filter(exp: str) -> Callable:
 88        """Filter a dataframe based on an expression.
 89
 90        Args:
 91            exp: filter expression.
 92
 93        Returns:
 94            A function to be called in .transform() spark function.
 95        """
 96
 97        def inner(df: DataFrame) -> DataFrame:
 98            return df.filter(exp)  # type: ignore
 99
100        return inner
101
102    @staticmethod
103    def column_filter_exp(exp: List[str]) -> Callable:
104        """Filter a dataframe's columns based on a list of SQL expressions.
105
106        Args:
107            exp: column filter expressions.
108
109        Returns:
110            A function to be called in .transform() spark function.
111        """
112
113        def inner(df: DataFrame) -> DataFrame:
114            return df.selectExpr(*exp)  # type: ignore
115
116        return inner
117
118    @staticmethod
119    def drop_duplicate_rows(
120        cols: List[str] = None, watermarker: dict = None
121    ) -> Callable:
122        """Drop duplicate rows using spark function dropDuplicates().
123
124        This transformer can be used with or without arguments.
125        The provided argument needs to be a list of columns.
126        For example: [“Name”,”VAT”] will drop duplicate records within
127        "Name" and "VAT" columns.
128        If the transformer is used without providing any columns list or providing
129        an empty list, such as [] the result will be the same as using
130        the distinct() pyspark function. If the watermark dict is present it will
131        ensure that the drop operation will apply to rows within the watermark timeline
132        window.
133
134
135        Args:
136            cols: column names.
137            watermarker: properties to apply watermarker to the transformer.
138
139        Returns:
140            A function to be called in .transform() spark function.
141        """
142
143        def inner(df: DataFrame) -> DataFrame:
144            if watermarker:
145                df = Watermarker.with_watermark(
146                    watermarker["col"], watermarker["watermarking_time"]
147                )(df)
148            if not cols:
149                return df.dropDuplicates()
150            else:
151                return df.dropDuplicates(cols)
152
153        return inner