Skip to content

Kafka writer

Module that defines the behaviour to write to Kafka.

KafkaWriter

Bases: Writer

Class to write to a Kafka target.

Source code in mkdocs/lakehouse_engine/packages/io/writers/kafka_writer.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
class KafkaWriter(Writer):
    """Class to write to a Kafka target."""

    def __init__(
        self, output_spec: OutputSpec, df: DataFrame, data: OrderedDict
    ) -> None:
        """Construct KafkaWriter instances.

        Args:
            output_spec: output specification.
            df: dataframe to be written.
            data: list of all dfs generated on previous steps before writer.
        """
        super().__init__(output_spec, df, data)

    def write(self) -> None:
        """Write data to Kafka."""
        if not self._df.isStreaming:
            self._write_to_kafka_in_batch_mode(self._df, self._output_spec)
        else:
            self._write_to_kafka_in_streaming_mode(
                self._df, self._output_spec, self._data
            )

    @staticmethod
    def _write_to_kafka_in_batch_mode(df: DataFrame, output_spec: OutputSpec) -> None:
        """Write to Kafka in batch mode.

        Args:
            df: dataframe to write.
            output_spec: output specification.
        """
        df.write.format(output_spec.data_format).options(
            **output_spec.options if output_spec.options else {}
        ).mode(output_spec.write_type).save()

    @staticmethod
    def _write_to_kafka_in_streaming_mode(
        df: DataFrame, output_spec: OutputSpec, data: OrderedDict
    ) -> None:
        """Write to kafka in streaming mode.

        Args:
            df: dataframe to write.
            output_spec: output specification.
            data: list of all dfs generated on previous steps before writer.
        """
        df_writer = df.writeStream.trigger(**Writer.get_streaming_trigger(output_spec))

        if (
            output_spec.streaming_micro_batch_transformers
            or output_spec.streaming_micro_batch_dq_processors
        ):
            stream_df = (
                df_writer.options(**output_spec.options if output_spec.options else {})
                .foreachBatch(
                    KafkaWriter._write_transformed_micro_batch(output_spec, data)
                )
                .start()
            )
        else:
            stream_df = (
                df_writer.format(output_spec.data_format)
                .options(**output_spec.options if output_spec.options else {})
                .start()
            )

        if output_spec.streaming_await_termination:
            stream_df.awaitTermination(output_spec.streaming_await_termination_timeout)

    @staticmethod
    def _write_transformed_micro_batch(
        output_spec: OutputSpec,
        data: OrderedDict | None,
    ) -> Callable:
        """Define how to write a streaming micro batch after transforming it.

        Args:
            output_spec: output specification.
            data: list of all dfs generated on previous steps before writer.

        Returns:
            A function to be executed in the foreachBatch spark write method.
        """

        def inner(batch_df: DataFrame, batch_id: int) -> None:
            ExecEnv.get_for_each_batch_session(batch_df)
            transformed_df = Writer.get_transformed_micro_batch(
                output_spec, batch_df, batch_id, data
            )

            if output_spec.streaming_micro_batch_dq_processors:
                transformed_df = Writer.run_micro_batch_dq_process(
                    transformed_df, output_spec.streaming_micro_batch_dq_processors
                )

            KafkaWriter._write_to_kafka_in_batch_mode(transformed_df, output_spec)

        return inner

__init__(output_spec, df, data)

Construct KafkaWriter instances.

Parameters:

Name Type Description Default
output_spec OutputSpec

output specification.

 required
df DataFrame

dataframe to be written.

 required
data OrderedDict

list of all dfs generated on previous steps before writer.

 required
Source code in mkdocs/lakehouse_engine/packages/io/writers/kafka_writer.py 
16
17
18
19
20
21
22
23
24
25
26
def __init__(
    self, output_spec: OutputSpec, df: DataFrame, data: OrderedDict
) -> None:
    """Construct KafkaWriter instances.

    Args:
        output_spec: output specification.
        df: dataframe to be written.
        data: list of all dfs generated on previous steps before writer.
    """
    super().__init__(output_spec, df, data)

write()

Write data to Kafka.

Source code in mkdocs/lakehouse_engine/packages/io/writers/kafka_writer.py 
28
29
30
31
32
33
34
35
def write(self) -> None:
    """Write data to Kafka."""
    if not self._df.isStreaming:
        self._write_to_kafka_in_batch_mode(self._df, self._output_spec)
    else:
        self._write_to_kafka_in_streaming_mode(
            self._df, self._output_spec, self._data
        )