lakehouse_engine.transformers.condensers

  1"""Condensers module."""
  2
  3from typing import Callable, List, Optional
  4
  5from pyspark.sql import DataFrame, Window
  6from pyspark.sql.functions import col, row_number
  7
  8from lakehouse_engine.transformers.exceptions import (
  9    UnsupportedStreamingTransformerException,
 10    WrongArgumentsException,
 11)
 12from lakehouse_engine.utils.logging_handler import LoggingHandler
 13
 14
 15class Condensers(object):
 16    """Class containing all the functions to condensate data for later merges."""
 17
 18    _logger = LoggingHandler(__name__).get_logger()
 19
 20    @classmethod
 21    def condense_record_mode_cdc(
 22        cls,
 23        business_key: List[str],
 24        record_mode_col: str,
 25        valid_record_modes: List[str],
 26        ranking_key_desc: Optional[List[str]] = None,
 27        ranking_key_asc: Optional[List[str]] = None,
 28    ) -> Callable:
 29        """Condense Change Data Capture (CDC) based on record_mode strategy.
 30
 31        This CDC data is particularly seen in some CDC enabled systems. Other systems
 32        may have different CDC strategies.
 33
 34        Args:
 35            business_key: The business key (logical primary key) of the data.
 36            ranking_key_desc: In this type of CDC condensation the data needs to be
 37                in descending order in a certain way, using columns specified in this
 38                parameter.
 39            ranking_key_asc: In this type of CDC condensation the data needs to be
 40                in ascending order in a certain way, using columns specified in
 41                this parameter.
 42            record_mode_col: Name of the record mode input_col.
 43            valid_record_modes: Depending on the context, not all record modes may be
 44                considered for condensation. Use this parameter to skip those.
 45
 46        Returns:
 47            A function to be executed in the .transform() spark function.
 48        """
 49        if not ranking_key_desc and not ranking_key_asc:
 50            raise WrongArgumentsException(
 51                "The condense_record_mode_cdc transformer requires data to be either"
 52                "in descending or ascending order, but no arguments for ordering"
 53                "were provided."
 54            )
 55
 56        def inner(df: DataFrame) -> DataFrame:
 57            if not df.isStreaming:
 58                partition_window = Window.partitionBy(
 59                    [col(c) for c in business_key]
 60                ).orderBy(
 61                    [
 62                        col(c).desc()
 63                        for c in (ranking_key_desc if ranking_key_desc else [])
 64                    ]  # type: ignore
 65                    + [
 66                        col(c).asc()
 67                        for c in (ranking_key_asc if ranking_key_asc else [])
 68                    ]  # type: ignore
 69                )
 70
 71                return (
 72                    df.withColumn("ranking", row_number().over(partition_window))
 73                    .filter(
 74                        col(record_mode_col).isNull()
 75                        | col(record_mode_col).isin(valid_record_modes)
 76                    )
 77                    .filter(col("ranking") == 1)
 78                    .drop("ranking")
 79                )
 80            else:
 81                raise UnsupportedStreamingTransformerException(
 82                    "Transformer condense_record_mode_cdc is not supported in "
 83                    "streaming mode."
 84                )
 85
 86        return inner
 87
 88    @classmethod
 89    def group_and_rank(
 90        cls, group_key: List[str], ranking_key: List[str], descending: bool = True
 91    ) -> Callable:
 92        """Condense data based on a simple group by + take latest mechanism.
 93
 94        Args:
 95            group_key: list of column names to use in the group by.
 96            ranking_key: the data needs to be in descending order using columns
 97                specified in this parameter.
 98            descending: if the ranking considers descending order or not. Defaults to
 99                True.
100
101        Returns:
102            A function to be executed in the .transform() spark function.
103        """
104
105        def inner(df: DataFrame) -> DataFrame:
106            if not df.isStreaming:
107                partition_window = Window.partitionBy(
108                    [col(c) for c in group_key]
109                ).orderBy(
110                    [
111                        col(c).desc() if descending else col(c).asc()
112                        for c in (ranking_key if ranking_key else [])
113                    ]  # type: ignore
114                )
115
116                return (
117                    df.withColumn("ranking", row_number().over(partition_window))
118                    .filter(col("ranking") == 1)
119                    .drop("ranking")
120                )
121            else:
122                raise UnsupportedStreamingTransformerException(
123                    "Transformer group_and_rank is not supported in streaming mode."
124                )
125
126        return inner