lakehouse_engine.transformers.null_handlers

 1"""Module with null handlers transformers."""
 2
 3from typing import Callable, List
 4
 5from pyspark.sql import DataFrame
 6
 7from lakehouse_engine.utils.logging_handler import LoggingHandler
 8
 9
10class NullHandlers(object):
11    """Class containing null handler transformers."""
12
13    _logger = LoggingHandler(__name__).get_logger()
14
15    @classmethod
16    def replace_nulls(
17        cls,
18        replace_on_nums: bool = True,
19        default_num_value: int = -999,
20        replace_on_strings: bool = True,
21        default_string_value: str = "UNKNOWN",
22        subset_cols: List[str] = None,
23    ) -> Callable:
24        """Replace nulls in a dataframe.
25
26        Args:
27            replace_on_nums: if it is to replace nulls on numeric columns.
28                Applies to ints, longs and floats.
29            default_num_value: default integer value to use as replacement.
30            replace_on_strings: if it is to replace nulls on string columns.
31            default_string_value: default string value to use as replacement.
32            subset_cols: list of columns in which to replace nulls. If not
33                provided, all nulls in all columns will be replaced as specified.
34
35        Returns:
36            A function to be called in .transform() spark function.
37        """
38
39        def inner(df: DataFrame) -> DataFrame:
40            if replace_on_nums:
41                df = df.na.fill(default_num_value, subset_cols)
42            if replace_on_strings:
43                df = df.na.fill(default_string_value, subset_cols)
44
45            return df
46
47        return inner