DataSourceStreamArrowWriter

Базовый класс для записи потоков данных, обрабатывающих данные с помощью PyArrow RecordBatch.

В отличие DataSourceStreamWriterот итератора объектов Spark Row , этот класс оптимизирован для формата стрелки при записи потоковых данных. Она может повысить производительность при взаимодействии с системами или библиотеками, которые изначально поддерживают стрелку для вариантов использования потоковой передачи. Реализуйте этот класс и верните экземпляр, DataSource.streamWriter() чтобы создать запись источника данных в качестве приемника потоковой передачи с помощью стрелки.

Синтаксис

from pyspark.sql.datasource import DataSourceStreamArrowWriter

class MyDataSourceStreamArrowWriter(DataSourceStreamArrowWriter):
    def write(self, iterator):
        ...

Методы

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

Примечания

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

Примеры

Реализуйте модуль записи потоков на основе стрелки, который подсчитывает строки на микробатч:

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

@dataclass
class MyCommitMessage(WriterCommitMessage):
    num_rows: int

class MyDataSourceStreamArrowWriter(DataSourceStreamArrowWriter):
    def write(self, iterator):
        total_rows = 0
        for batch in iterator:
            total_rows += len(batch)
        return MyCommitMessage(num_rows=total_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")