DataSourceStreamWriter

Базовый класс для записи потоков данных.

Записи потоков данных отвечают за запись данных в приемник потоковой передачи. Реализуйте этот класс и верните экземпляр для DataSource.streamWriter() записи источника данных в качестве приемника потоковой передачи. write() вызывается исполнителями для каждого микробатча и commit()abort() вызывается на драйвере после завершения всех задач в микробатче.

Синтаксис

from pyspark.sql.datasource import DataSourceStreamWriter

class MyDataSourceStreamWriter(DataSourceStreamWriter):
    def write(self, iterator):
        ...

Методы

Метод Описание
write(iterator) Записывает данные в приемник потоковой передачи. Звонил исполнителям один раз на микробатч. Принимает итератор объектов и возвращает Rowсообщение о фиксацииWriterCommitMessage.None Этот метод является абстрактным и должен быть реализован.
commit(messages, batchId) Фиксирует микробатч с помощью списка сообщений фиксации, собранных всеми исполнителями. Вызывается на драйвере, когда все задачи в микробатче успешно выполняются.
abort(messages, batchId) Прерывает микробатч с помощью списка сообщений фиксации, собранных со всех исполнителей. Вызывается на драйвере, когда произошел сбой одной или нескольких задач в микробатче.

Примечания

  • Драйвер собирает сообщения фиксации от всех исполнителей и передает их commit() в случае успешного выполнения всех задач или в abort() случае сбоя любой задачи.
  • Если задача записи завершается ошибкой, сообщение о фиксации будет в None списке, переданном commit() или abort().
  • batchId уникально идентифицирует каждый микробатч и увеличивается на 1 с каждым обработанным микробатчем.

Примеры

Реализуйте модуль записи потоков, добавляющий строки в файл:

from dataclasses import dataclass
from pyspark.sql.datasource import DataSource, DataSourceStreamWriter, WriterCommitMessage

@dataclass
class MyCommitMessage(WriterCommitMessage):
    num_rows: int

class MyDataSourceStreamWriter(DataSourceStreamWriter):
    def __init__(self, options):
        self.path = options.get("path")

    def write(self, iterator):
        rows = list(iterator)
        with open(self.path, "a") as f:
            for row in rows:
                f.write(str(row) + "\n")
        return MyCommitMessage(num_rows=len(rows))

    def commit(self, messages, batchId):
        total = sum(m.num_rows for m in messages if m is not None)
        print(f"Committed batch {batchId} with {total} rows")

    def abort(self, messages, batchId):
        print(f"Batch {batchId} failed, performing cleanup")