Skip to content

File writer

Module to define behaviour to write to files.

FileWriter

Bases: Writer

Class to write data to files.

Source code in mkdocs/lakehouse_engine/packages/io/writers/file_writer.py 
 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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
class FileWriter(Writer):
    """Class to write data to files."""

    def __init__(
        self, output_spec: OutputSpec, df: DataFrame, data: OrderedDict | None
    ) -> None:
        """Construct FileWriter 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 files."""
        if not self._df.isStreaming:
            self._write_to_files_in_batch_mode(self._df, self._output_spec)
        else:
            self._write_to_files_in_streaming_mode(
                self._df, self._output_spec, self._data
            )

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

        Args:
            df: dataframe to write.
            output_spec: output specification.
        """
        if output_spec.data_format == OutputFormat.PAIMON.value:
            PaimonUtils.ensure_table_exists(ExecEnv.SESSION, df, output_spec)
            df = PaimonUtils.align_dataframe_with_table_schema(
                ExecEnv.SESSION, df, output_spec
            )

        writer = df.write.format(output_spec.data_format)
        if output_spec.partitions:
            writer = writer.partitionBy(*output_spec.partitions)
        writer.options(**output_spec.options if output_spec.options else {}).mode(
            output_spec.write_type
        ).save(output_spec.location)

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

        Args:
            df: dataframe to write.
            output_spec: output specification.
            data: list of all dfs generated on previous steps before writer.
        """
        if output_spec.data_format == OutputFormat.PAIMON.value:
            PaimonUtils.ensure_table_exists(ExecEnv.SESSION, df, output_spec)
            df = PaimonUtils.align_dataframe_with_table_schema(
                ExecEnv.SESSION, df, output_spec
            )

        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(
                    FileWriter._write_transformed_micro_batch(output_spec, data)
                )
                .start()
            )
        else:
            stream_writer = df_writer.format(output_spec.data_format)
            if output_spec.partitions:
                stream_writer = stream_writer.partitionBy(*output_spec.partitions)
            stream_df = (
                stream_writer.options(
                    **output_spec.options if output_spec.options else {}
                )
                .outputMode(output_spec.write_type)
                .start(output_spec.location)
            )

        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
                )

            FileWriter._write_to_files_in_batch_mode(transformed_df, output_spec)

        return inner

__init__(output_spec, df, data)

Construct FileWriter instances.

Parameters:

Name Type Description Default
output_spec OutputSpec

output specification

 required
df DataFrame

dataframe to be written.

 required
data OrderedDict | None

list of all dfs generated on previous steps before writer.

 required
Source code in mkdocs/lakehouse_engine/packages/io/writers/file_writer.py 
17
18
19
20
21
22
23
24
25
26
27
def __init__(
    self, output_spec: OutputSpec, df: DataFrame, data: OrderedDict | None
) -> None:
    """Construct FileWriter 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 files.

Source code in mkdocs/lakehouse_engine/packages/io/writers/file_writer.py 
29
30
31
32
33
34
35
36
def write(self) -> None:
    """Write data to files."""
    if not self._df.isStreaming:
        self._write_to_files_in_batch_mode(self._df, self._output_spec)
    else:
        self._write_to_files_in_streaming_mode(
            self._df, self._output_spec, self._data
        )