lakehouse_engine.io.writers.table_writer

  1"""Module that defines the behaviour to write to tables."""
  2
  3from typing import Any, Callable, OrderedDict
  4
  5from pyspark.sql import DataFrame
  6
  7from lakehouse_engine.core.definitions import OutputFormat, OutputSpec
  8from lakehouse_engine.core.exec_env import ExecEnv
  9from lakehouse_engine.io.writer import Writer
 10
 11
 12class TableWriter(Writer):
 13    """Class to write to a table."""
 14
 15    def __init__(self, output_spec: OutputSpec, df: DataFrame, data: OrderedDict):
 16        """Construct TableWriter instances.
 17
 18        Args:
 19            output_spec: output specification.
 20            df: dataframe to be written.
 21            data: list of all dfs generated on previous steps before writer.
 22        """
 23        super().__init__(output_spec, df, data)
 24
 25    def write(self) -> None:
 26        """Write data to a table.
 27
 28        After the write operation we repair the table (e.g., update partitions).
 29        However, there's a caveat to this, which is the fact that this repair
 30        operation is not reachable if we are running long-running streaming mode.
 31        Therefore, we recommend not using the TableWriter with formats other than
 32        delta lake for those scenarios (as delta lake does not need msck repair).
 33        So, you can: 1) use delta lake format for the table; 2) use the FileWriter
 34        and run the repair with a certain frequency in a separate task of your
 35        pipeline.
 36        """
 37        if not self._df.isStreaming:
 38            self._write_to_table_in_batch_mode(self._df, self._output_spec)
 39        else:
 40            df_writer = self._df.writeStream.trigger(
 41                **Writer.get_streaming_trigger(self._output_spec)
 42            )
 43
 44            if (
 45                self._output_spec.streaming_micro_batch_transformers
 46                or self._output_spec.streaming_micro_batch_dq_processors
 47            ):
 48                stream_df = (
 49                    df_writer.options(
 50                        **self._output_spec.options if self._output_spec.options else {}
 51                    )
 52                    .foreachBatch(
 53                        self._write_transformed_micro_batch(
 54                            self._output_spec, self._data
 55                        )
 56                    )
 57                    .start()
 58                )
 59
 60                if self._output_spec.streaming_await_termination:
 61                    stream_df.awaitTermination(
 62                        self._output_spec.streaming_await_termination_timeout
 63                    )
 64            else:
 65                self._write_to_table_in_streaming_mode(df_writer, self._output_spec)
 66
 67        if (
 68            self._output_spec.data_format != OutputFormat.DELTAFILES.value
 69            and self._output_spec.partitions
 70        ):
 71            ExecEnv.SESSION.sql(f"MSCK REPAIR TABLE {self._output_spec.db_table}")
 72
 73    @staticmethod
 74    def _write_to_table_in_batch_mode(df: DataFrame, output_spec: OutputSpec) -> None:
 75        """Write to a metastore table in batch mode.
 76
 77        Args:
 78            df: dataframe to write.
 79            output_spec: output specification.
 80        """
 81        df_writer = df.write.format(output_spec.data_format)
 82
 83        if output_spec.partitions:
 84            df_writer = df_writer.partitionBy(output_spec.partitions)
 85
 86        if output_spec.location:
 87            df_writer = df_writer.options(
 88                path=output_spec.location,
 89                **output_spec.options if output_spec.options else {},
 90            )
 91        else:
 92            df_writer = df_writer.options(
 93                **output_spec.options if output_spec.options else {}
 94            )
 95
 96        df_writer.mode(output_spec.write_type).saveAsTable(output_spec.db_table)
 97
 98    @staticmethod
 99    def _write_to_table_in_streaming_mode(
100        df_writer: Any, output_spec: OutputSpec
101    ) -> None:
102        """Write to a metastore table in streaming mode.
103
104        Args:
105            df_writer: dataframe writer.
106            output_spec: output specification.
107        """
108        df_writer = df_writer.outputMode(output_spec.write_type).format(
109            output_spec.data_format
110        )
111
112        if output_spec.partitions:
113            df_writer = df_writer.partitionBy(output_spec.partitions)
114
115        if output_spec.location:
116            df_writer = df_writer.options(
117                path=output_spec.location,
118                **output_spec.options if output_spec.options else {},
119            )
120        else:
121            df_writer = df_writer.options(
122                **output_spec.options if output_spec.options else {}
123            )
124
125        if output_spec.streaming_await_termination:
126            df_writer.toTable(output_spec.db_table).awaitTermination(
127                output_spec.streaming_await_termination_timeout
128            )
129        else:
130            df_writer.toTable(output_spec.db_table)
131
132    @staticmethod
133    def _write_transformed_micro_batch(  # type: ignore
134        output_spec: OutputSpec, data: OrderedDict
135    ) -> Callable:
136        """Define how to write a streaming micro batch after transforming it.
137
138        Args:
139            output_spec: output specification.
140            data: list of all dfs generated on previous steps before writer.
141
142        Returns:
143            A function to be executed in the foreachBatch spark write method.
144        """
145
146        def inner(batch_df: DataFrame, batch_id: int) -> None:
147            transformed_df = Writer.get_transformed_micro_batch(
148                output_spec, batch_df, batch_id, data
149            )
150
151            if output_spec.streaming_micro_batch_dq_processors:
152                transformed_df = Writer.run_micro_batch_dq_process(
153                    transformed_df, output_spec.streaming_micro_batch_dq_processors
154                )
155
156            TableWriter._write_to_table_in_batch_mode(transformed_df, output_spec)
157
158        return inner