lakehouse_engine.terminators.dataset_optimizer

  1"""Module with dataset optimizer terminator."""
  2
  3from typing import List, Optional
  4
  5from pyspark.sql.utils import AnalysisException, ParseException
  6
  7from lakehouse_engine.core.table_manager import TableManager
  8from lakehouse_engine.transformers.exceptions import WrongArgumentsException
  9from lakehouse_engine.utils.logging_handler import LoggingHandler
 10
 11
 12class DatasetOptimizer(object):
 13    """Class with dataset optimizer terminator."""
 14
 15    _logger = LoggingHandler(__name__).get_logger()
 16
 17    @classmethod
 18    def optimize_dataset(
 19        cls,
 20        db_table: Optional[str] = None,
 21        location: Optional[str] = None,
 22        compute_table_stats: bool = True,
 23        vacuum: bool = True,
 24        vacuum_hours: int = 720,
 25        optimize: bool = True,
 26        optimize_where: Optional[str] = None,
 27        optimize_zorder_col_list: Optional[List[str]] = None,
 28        debug: bool = False,
 29    ) -> None:
 30        """Optimize a dataset based on a set of pre-conceived optimizations.
 31
 32        Most of the time the dataset is a table, but it can be a file-based one only.
 33
 34        Args:
 35            db_table: `database_name.table_name`.
 36            location: dataset/table filesystem location.
 37            compute_table_stats: to compute table statistics or not.
 38            vacuum: (delta lake tables only) whether to vacuum the delta lake
 39                table or not.
 40            vacuum_hours: (delta lake tables only) number of hours to consider
 41                in vacuum operation.
 42            optimize: (delta lake tables only) whether to optimize the table or
 43                not. Custom optimize parameters can be supplied through ExecEnv (Spark)
 44                configs
 45            optimize_where: expression to use in the optimize function.
 46            optimize_zorder_col_list: (delta lake tables only) list of
 47                columns to consider in the zorder optimization process. Custom optimize
 48                parameters can be supplied through ExecEnv (Spark) configs.
 49            debug: flag indicating if we are just debugging this for local
 50                tests and therefore pass through all the exceptions to perform some
 51                assertions in local tests.
 52        """
 53        if optimize:
 54            if debug:
 55                try:
 56                    cls._optimize(
 57                        db_table, location, optimize_where, optimize_zorder_col_list
 58                    )
 59                except ParseException:
 60                    pass
 61            else:
 62                cls._optimize(
 63                    db_table, location, optimize_where, optimize_zorder_col_list
 64                )
 65
 66        if vacuum:
 67            cls._vacuum(db_table, location, vacuum_hours)
 68
 69        if compute_table_stats:
 70            if debug:
 71                try:
 72                    cls._compute_table_stats(db_table)
 73                except AnalysisException:
 74                    pass
 75            else:
 76                cls._compute_table_stats(db_table)
 77
 78    @classmethod
 79    def _compute_table_stats(cls, db_table: str) -> None:
 80        """Compute table statistics.
 81
 82        Args:
 83            db_table: `<db>.<table>` string.
 84        """
 85        if not db_table:
 86            raise WrongArgumentsException("A table needs to be provided.")
 87
 88        config = {"function": "compute_table_statistics", "table_or_view": db_table}
 89        cls._logger.info(f"Computing table statistics for {db_table}...")
 90        TableManager(config).compute_table_statistics()
 91
 92    @classmethod
 93    def _vacuum(cls, db_table: str, location: str, hours: int) -> None:
 94        """Vacuum a delta table.
 95
 96        Args:
 97            db_table: `<db>.<table>` string. Takes precedence over location.
 98            location: location of the delta table.
 99            hours: number of hours to consider in vacuum operation.
100        """
101        if not db_table and not location:
102            raise WrongArgumentsException("A table or location need to be provided.")
103
104        table_or_location = db_table if db_table else f"delta.`{location}`"
105
106        config = {
107            "function": "compute_table_statistics",
108            "table_or_view": table_or_location,
109            "vacuum_hours": hours,
110        }
111        cls._logger.info(f"Vacuuming table {table_or_location}...")
112        TableManager(config).vacuum()
113
114    @classmethod
115    def _optimize(
116        cls, db_table: str, location: str, where: str, zorder_cols: List[str]
117    ) -> None:
118        """Optimize a delta table.
119
120        Args:
121            db_table: `<db>.<table>` string. Takes precedence over location.
122            location: location of the delta table.
123            where: expression to use in the optimize function.
124            zorder_cols: list of columns to consider in the zorder optimization process.
125        """
126        if not db_table and not location:
127            raise WrongArgumentsException("A table or location needs to be provided.")
128
129        table_or_location = db_table if db_table else f"delta.`{location}`"
130
131        config = {
132            "function": "compute_table_statistics",
133            "table_or_view": table_or_location,
134            "optimize_where": where,
135            "optimize_zorder_col_list": ",".join(zorder_cols if zorder_cols else []),
136        }
137        cls._logger.info(f"Optimizing table {table_or_location}...")
138        TableManager(config).optimize()