lakehouse_engine.transformers.column_creators

  1"""Column creators transformers module."""
  2
  3from typing import Any, Callable, Dict
  4
  5from pyspark.sql import DataFrame, Window
  6from pyspark.sql.functions import col, lit, monotonically_increasing_id, row_number
  7from pyspark.sql.types import IntegerType
  8
  9from lakehouse_engine.transformers.exceptions import (
 10    UnsupportedStreamingTransformerException,
 11)
 12from lakehouse_engine.utils.logging_handler import LoggingHandler
 13
 14
 15class ColumnCreators(object):
 16    """Class containing all functions that can create columns to add value."""
 17
 18    _logger = LoggingHandler(__name__).get_logger()
 19
 20    @classmethod
 21    def with_row_id(
 22        cls,
 23        output_col: str = "lhe_row_id",
 24    ) -> Callable:
 25        """Create a sequential but not consecutive id.
 26
 27        Args:
 28            output_col: optional name of the output column.
 29
 30        Returns:
 31            A function to be executed in the .transform() spark function.
 32        """
 33
 34        def inner(df: DataFrame) -> DataFrame:
 35            if not df.isStreaming:
 36                return df.withColumn(output_col, monotonically_increasing_id())
 37            else:
 38                raise UnsupportedStreamingTransformerException(
 39                    "Transformer with_row_id is not supported in streaming mode."
 40                )
 41
 42        return inner
 43
 44    @classmethod
 45    def with_auto_increment_id(
 46        cls, output_col: str = "lhe_row_id", rdd: bool = True
 47    ) -> Callable:
 48        """Create a sequential and consecutive id.
 49
 50        Args:
 51            output_col: optional name of the output column.
 52            rdd: optional parameter to use spark rdd.
 53
 54        Returns:
 55            A function to be executed in the .transform() spark function.
 56        """
 57
 58        def inner(df: DataFrame) -> DataFrame:
 59            if not df.isStreaming:
 60                if len(df.take(1)) == 0:
 61                    # if df is empty we have to prevent the algorithm from failing
 62                    return df.withColumn(output_col, lit(None).cast(IntegerType()))
 63                elif rdd:
 64                    return (
 65                        df.rdd.zipWithIndex()
 66                        .toDF()
 67                        .select(col("_1.*"), col("_2").alias(output_col))
 68                    )
 69                else:
 70                    w = Window.orderBy(monotonically_increasing_id())
 71                    return df.withColumn(output_col, (row_number().over(w)) - 1)
 72
 73            else:
 74                raise UnsupportedStreamingTransformerException(
 75                    "Transformer with_auto_increment_id is not supported in "
 76                    "streaming mode."
 77                )
 78
 79        return inner
 80
 81    @classmethod
 82    def with_literals(
 83        cls,
 84        literals: Dict[str, Any],
 85    ) -> Callable:
 86        """Create columns given a map of column names and literal values (constants).
 87
 88        Args:
 89            Dict[str, Any] literals: map of column names and literal values (constants).
 90
 91        Returns:
 92            Callable: A function to be executed in the .transform() spark function.
 93        """
 94
 95        def inner(df: DataFrame) -> DataFrame:
 96            df_with_literals = df
 97            for name, value in literals.items():
 98                df_with_literals = df_with_literals.withColumn(name, lit(value))
 99            return df_with_literals
100
101        return inner