lakehouse_engine.io.writers.kafka_writer

  1"""Module that defines the behaviour to write to Kafka."""
  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
  9
 10
 11class KafkaWriter(Writer):
 12    """Class to write to a Kafka target."""
 13
 14    def __init__(self, output_spec: OutputSpec, df: DataFrame, data: OrderedDict):
 15        """Construct KafkaWriter instances.
 16
 17        Args:
 18            output_spec: output specification.
 19            df: dataframe to be written.
 20            data: list of all dfs generated on previous steps before writer.
 21        """
 22        super().__init__(output_spec, df, data)
 23
 24    def write(self) -> None:
 25        """Write data to Kafka."""
 26        if not self._df.isStreaming:
 27            self._write_to_kafka_in_batch_mode(self._df, self._output_spec)
 28        else:
 29            self._write_to_kafka_in_streaming_mode(
 30                self._df, self._output_spec, self._data
 31            )
 32
 33    @staticmethod
 34    def _write_to_kafka_in_batch_mode(df: DataFrame, output_spec: OutputSpec) -> None:
 35        """Write to Kafka in batch mode.
 36
 37        Args:
 38            df: dataframe to write.
 39            output_spec: output specification.
 40        """
 41        df.write.format(output_spec.data_format).options(
 42            **output_spec.options if output_spec.options else {}
 43        ).mode(output_spec.write_type).save()
 44
 45    @staticmethod
 46    def _write_to_kafka_in_streaming_mode(
 47        df: DataFrame, output_spec: OutputSpec, data: OrderedDict
 48    ) -> None:
 49        """Write to kafka in streaming mode.
 50
 51        Args:
 52            df: dataframe to write.
 53            output_spec: output specification.
 54            data: list of all dfs generated on previous steps before writer.
 55        """
 56        df_writer = df.writeStream.trigger(**Writer.get_streaming_trigger(output_spec))
 57
 58        if (
 59            output_spec.streaming_micro_batch_transformers
 60            or output_spec.streaming_micro_batch_dq_processors
 61        ):
 62            stream_df = (
 63                df_writer.options(**output_spec.options if output_spec.options else {})
 64                .foreachBatch(
 65                    KafkaWriter._write_transformed_micro_batch(output_spec, data)
 66                )
 67                .start()
 68            )
 69        else:
 70            stream_df = (
 71                df_writer.format(output_spec.data_format)
 72                .options(**output_spec.options if output_spec.options else {})
 73                .start()
 74            )
 75
 76        if output_spec.streaming_await_termination:
 77            stream_df.awaitTermination(output_spec.streaming_await_termination_timeout)
 78
 79    @staticmethod
 80    def _write_transformed_micro_batch(  # type: ignore
 81        output_spec: OutputSpec, data: OrderedDict
 82    ) -> Callable:
 83        """Define how to write a streaming micro batch after transforming it.
 84
 85        Args:
 86            output_spec: output specification.
 87            data: list of all dfs generated on previous steps before writer.
 88
 89        Returns:
 90            A function to be executed in the foreachBatch spark write method.
 91        """
 92
 93        def inner(batch_df: DataFrame, batch_id: int) -> None:
 94            transformed_df = Writer.get_transformed_micro_batch(
 95                output_spec, batch_df, batch_id, data
 96            )
 97
 98            if output_spec.streaming_micro_batch_dq_processors:
 99                transformed_df = Writer.run_micro_batch_dq_process(
100                    transformed_df, output_spec.streaming_micro_batch_dq_processors
101                )
102
103            KafkaWriter._write_to_kafka_in_batch_mode(transformed_df, output_spec)
104
105        return inner