DataSourceArrowWriter

Základní třída pro zapisovače zdrojů dat, které zpracovávají data pomocí PyArrow's RecordBatch.

Na rozdíl od DataSourceWritertoho, který pracuje s iterátorem objektů Spark Row , je tato třída optimalizována pro formát Šipka při zápisu dat. Může nabídnout lepší výkon při spolupráci se systémy nebo knihovnami, které nativně podporují šipku. Implementujte tuto třídu a vraťte instanci, DataSource.writer() aby byl zdroj dat zapisovatelný pomocí šipky.

Syntaxe

from pyspark.sql.datasource import DataSourceArrowWriter

class MyDataSourceArrowWriter(DataSourceArrowWriter):
    def write(self, iterator):
        ...

Methods

Metoda Description
write(iterator) Zapíše iterátor objektů PyArrow RecordBatch do jímky. Volá se jednou na každém exekutoru. WriterCommitMessageVrátí zprávu , nebo None pokud neexistuje žádná potvrzovací zpráva. Tato metoda je abstraktní a musí být implementována.
commit(messages) Potvrdí úlohu zápisu pomocí seznamu zpráv potvrzení shromážděných ze všech exekutorů. Vyvoláno na ovladači, když se všechny úlohy úspěšně spustí. Zděděno z DataSourceWriter.
abort(messages) Přeruší úlohu zápisu pomocí seznamu zpráv potvrzení shromážděných ze všech exekutorů. Vyvoláno na ovladači, když došlo k selhání jedné nebo více úloh. Zděděno z DataSourceWriter.

Poznámky

  • Ovladač shromažďuje zprávy potvrzení ze všech exekutorů a předává je, commit() pokud jsou všechny úkoly úspěšné nebo pokud abort() některý úkol selže.
  • Pokud úloha zápisu selže, zpráva potvrzení bude None v seznamu předána commit() nebo abort().

Příklady

Implementujte zapisovač založený na šipkách, který počítá řádky napříč všemi dávkami:

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

@dataclass
class MyCommitMessage(WriterCommitMessage):
    num_rows: int

class MyDataSourceArrowWriter(DataSourceArrowWriter):
    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):
        total = sum(m.num_rows for m in messages if m is not None)
        print(f"Committed {total} rows")

    def abort(self, messages):
        print("Write job failed, performing cleanup")