lakehouse_engine.io.writers.console_writer

  1"""Module to define behaviour to write to console."""
  2
  3from typing import Callable, OrderedDict
  4
  5from pyspark.sql import DataFrame
  6
  7from lakehouse_engine.core.definitions import OutputSpec
  8from lakehouse_engine.io.writer import Writer
  9from lakehouse_engine.utils.logging_handler import LoggingHandler
 10
 11
 12class ConsoleWriter(Writer):
 13    """Class to write data to console."""
 14
 15    _logger = LoggingHandler(__name__).get_logger()
 16
 17    def __init__(self, output_spec: OutputSpec, df: DataFrame, data: OrderedDict):
 18        """Construct ConsoleWriter instances.
 19
 20        Args:
 21            output_spec: output specification
 22            df: dataframe to be written.
 23            data: list of all dfs generated on previous steps before writer.
 24        """
 25        super().__init__(output_spec, df, data)
 26
 27    def write(self) -> None:
 28        """Write data to console."""
 29        self._output_spec.options = (
 30            self._output_spec.options if self._output_spec.options else {}
 31        )
 32        if not self._df.isStreaming:
 33            self._logger.info("Dataframe preview:")
 34            self._show_df(self._df, self._output_spec)
 35        else:
 36            self._logger.info("Stream Dataframe preview:")
 37            self._write_to_console_in_streaming_mode(
 38                self._df, self._output_spec, self._data
 39            )
 40
 41    @staticmethod
 42    def _show_df(df: DataFrame, output_spec: OutputSpec) -> None:
 43        """Given a dataframe it applies Spark's show function to show it.
 44
 45        Args:
 46            df: dataframe to be shown.
 47            output_spec: output specification.
 48        """
 49        df.show(
 50            n=output_spec.options.get("limit", 20),
 51            truncate=output_spec.options.get("truncate", True),
 52            vertical=output_spec.options.get("vertical", False),
 53        )
 54
 55    @staticmethod
 56    def _show_streaming_df(output_spec: OutputSpec) -> Callable:
 57        """Define how to show a streaming df.
 58
 59        Args:
 60            output_spec: output specification.
 61
 62        Returns:
 63            A function to show df in the foreachBatch spark write method.
 64        """
 65
 66        def inner(batch_df: DataFrame, batch_id: int) -> None:
 67            ConsoleWriter._logger.info(f"Showing DF for batch {batch_id}")
 68            ConsoleWriter._show_df(batch_df, output_spec)
 69
 70        return inner
 71
 72    @staticmethod
 73    def _write_to_console_in_streaming_mode(
 74        df: DataFrame, output_spec: OutputSpec, data: OrderedDict
 75    ) -> None:
 76        """Write to console in streaming mode.
 77
 78        Args:
 79            df: dataframe to write.
 80            output_spec: output specification.
 81            data: list of all dfs generated on previous steps before writer.
 82        """
 83        df_writer = df.writeStream.trigger(**Writer.get_streaming_trigger(output_spec))
 84
 85        if (
 86            output_spec.streaming_micro_batch_transformers
 87            or output_spec.streaming_micro_batch_dq_processors
 88        ):
 89            stream_df = df_writer.foreachBatch(
 90                ConsoleWriter._write_transformed_micro_batch(output_spec, data)
 91            ).start()
 92        else:
 93            stream_df = df_writer.foreachBatch(
 94                ConsoleWriter._show_streaming_df(output_spec)
 95            ).start()
 96
 97        if output_spec.streaming_await_termination:
 98            stream_df.awaitTermination(output_spec.streaming_await_termination_timeout)
 99
100    @staticmethod
101    def _write_transformed_micro_batch(  # type: ignore
102        output_spec: OutputSpec, data: OrderedDict
103    ) -> Callable:
104        """Define how to write a streaming micro batch after transforming it.
105
106        Args:
107            output_spec: output specification.
108            data: list of all dfs generated on previous steps before writer.
109
110        Returns:
111            A function to be executed in the foreachBatch spark write method.
112        """
113
114        def inner(batch_df: DataFrame, batch_id: int) -> None:
115            transformed_df = Writer.get_transformed_micro_batch(
116                output_spec, batch_df, batch_id, data
117            )
118
119            if output_spec.streaming_micro_batch_dq_processors:
120                transformed_df = Writer.run_micro_batch_dq_process(
121                    transformed_df, output_spec.streaming_micro_batch_dq_processors
122                )
123
124            ConsoleWriter._show_df(transformed_df, output_spec)
125
126        return inner