Vlastní zdroje dat PySpark

Důležité

Vlastní zdroje dat PySpark jsou ve verzi Public Preview v Databricks Runtime 15.2 a vyšší. Podpora streamování je dostupná ve službě Databricks Runtime 15.3 a vyšší.

Rozhraní API zdroje dat Pythonu (PySpark) vytváří PySpark DataSource API, které umožňuje čtení z vlastních zdrojů dat a zápis do vlastních jímek dat v Apache Sparku pomocí Pythonu. Vlastní zdroje dat PySpark můžete použít k definování vlastních připojení k datovým systémům a implementaci dalších funkcí k vytváření opakovaně použitelných zdrojů dat.

DataSource – třída

PySpark DataSource je základní třída, která poskytuje metody pro vytváření čteček dat a zapisovačů dat.

Implementace podtřídy zdroje dat

V závislosti na vašem případu použití musí každá podtřída implementovat následující podtřídu, aby byl zdroj dat čitelný, zapisovatelný nebo obojí:

Vlastnost nebo metoda	Popis
name	Požadováno. Název zdroje dat
schema	Požadováno. Schéma zdroje dat ke čtení nebo zápisu
reader()	Musí se vrátit DataSourceReader , aby byl zdroj dat čitelný (dávka).
writer()	Musí se vrátit DataSourceWriter , aby byla datová jímka zapisovatelná (dávka).
streamReader() nebo simpleStreamReader()	Musí se vrátit DataSourceStreamReader , aby datový proud byl čitelný (streamování).
streamWriter()	Musí se vrátit, DataSourceStreamWriter aby datový proud byl zapisovatelný (streamování).

Poznámka:

Uživatelem definovaný DataSource, , DataSourceReader, DataSourceWriterDataSourceStreamReader, DataSourceStreamWritera jejich metody musí být schopni serializovat. Jinými slovy, musí se jednat o slovník nebo vnořený slovník, který obsahuje primitivní typ.

Registrace zdroje dat

Po implementaci rozhraní ho musíte zaregistrovat, pak ho můžete načíst nebo jinak použít, jak je znázorněno v následujícím příkladu:

# Register the data source
spark.dataSource.register(MyDataSourceClass)

# Read from a custom data source
spark.read.format("my_datasource_name").load().show()

Příklad 1: Vytvoření zdroje dat PySpark pro dávkové dotazy

Pokud chcete předvést možnosti čtečky dat PySpark DataSource, vytvořte zdroj dat, který generuje ukázková data pomocí faker balíčku Pythonu. Další informace najdete v fakerdokumentaci k fakeru.

faker Nainstalujte balíček pomocí následujícího příkazu:

%pip install faker

Krok 1: Definování ukázkového zdroje dat

Nejprve definujte nový zdroj dat PySpark jako podtřídu DataSource s názvem, schématem a čtenářem. Metoda musí být definována reader() pro čtení ze zdroje dat v dávkovém dotazu.

from pyspark.sql.datasource import DataSource, DataSourceReader
from pyspark.sql.types import StructType

class FakeDataSource(DataSource):
    """
    An example data source for batch query using the `faker` library.
    """

    @classmethod
    def name(cls):
        return "fake"

    def schema(self):
        return "name string, date string, zipcode string, state string"

    def reader(self, schema: StructType):
        return FakeDataSourceReader(schema, self.options)

Krok 2: Implementace čtečky dávkového dotazu

Dále implementujte logiku čtenáře, která vygeneruje ukázková data. K naplnění každého pole ve schématu použijte nainstalovanou faker knihovnu.

class FakeDataSourceReader(DataSourceReader):

    def __init__(self, schema, options):
        self.schema: StructType = schema
        self.options = options

    def read(self, partition):
        # Library imports must be within the method.
        from faker import Faker
        fake = Faker()

        # Every value in this `self.options` dictionary is a string.
        num_rows = int(self.options.get("numRows", 3))
        for _ in range(num_rows):
            row = []
            for field in self.schema.fields:
                value = getattr(fake, field.name)()
                row.append(value)
            yield tuple(row)

Krok 3: Registrace a použití ukázkového zdroje dat

Pokud chcete zdroj dat použít, zaregistrujte ho. Ve výchozím nastavení má FakeDataSource tři řádky a schéma obsahuje tato string pole: name, date, zipcode, state. Následující příklad zaregistruje, načte a vypíše ukázkový zdroj dat s výchozími nastaveními:

spark.dataSource.register(FakeDataSource)
spark.read.format("fake").load().show()

+-----------------+----------+-------+----------+
|             name|      date|zipcode|     state|
+-----------------+----------+-------+----------+
|Christine Sampson|1979-04-24|  79766|  Colorado|
|       Shelby Cox|2011-08-05|  24596|   Florida|
|  Amanda Robinson|2019-01-06|  57395|Washington|
+-----------------+----------+-------+----------+

Podporují se pouze string pole, ale můžete zadat schéma s libovolnými poli, která odpovídají faker polím poskytovatelů balíčků, aby se generovala náhodná data pro testování a vývoj. Následující příklad načte zdroj dat s name poli a company pole:

spark.read.format("fake").schema("name string, company string").load().show()

+---------------------+--------------+
|name                 |company       |
+---------------------+--------------+
|Tanner Brennan       |Adams Group   |
|Leslie Maxwell       |Santiago Group|
|Mrs. Jacqueline Brown|Maynard Inc   |
+---------------------+--------------+

Pokud chcete načíst zdroj dat s vlastním počtem řádků, zadejte numRows možnost. Následující příklad určuje 5 řádků:

spark.read.format("fake").option("numRows", 5).load().show()

+--------------+----------+-------+------------+
|          name|      date|zipcode|       state|
+--------------+----------+-------+------------+
|  Pam Mitchell|1988-10-20|  23788|   Tennessee|
|Melissa Turner|1996-06-14|  30851|      Nevada|
|  Brian Ramsey|2021-08-21|  55277|  Washington|
|  Caitlin Reed|1983-06-22|  89813|Pennsylvania|
| Douglas James|2007-01-18|  46226|     Alabama|
+--------------+----------+-------+------------+

Příklad 2: Vytvoření zdroje dat PySpark pro čtení a zápis streamování

Pokud chcete předvést možnosti čtečky datových proudů a zapisovače PySpark, vytvořte ukázkový zdroj dat, který vygeneruje dva řádky v každém mikrobatchu pomocí balíčku Pythonu faker . Další informace najdete v fakerdokumentaci k fakeru.

faker Nainstalujte balíček pomocí následujícího příkazu:

%pip install faker

Krok 1: Definování ukázkového zdroje dat

Nejprve definujte nový zdroj dat PySpark jako podtřídu DataSource s názvem, schématem a metodami streamReader() a streamWriter().

from pyspark.sql.datasource import DataSource, DataSourceStreamReader, SimpleDataSourceStreamReader, DataSourceStreamWriter
from pyspark.sql.types import StructType

class FakeStreamDataSource(DataSource):
    """
    An example data source for streaming read and write using the `faker` library.
    """

    @classmethod
    def name(cls):
        return "fakestream"

    def schema(self):
        return "name string, state string"

    def streamReader(self, schema: StructType):
        return FakeStreamReader(schema, self.options)

    # If you don't need partitioning, you can implement the simpleStreamReader method instead of streamReader.
    # def simpleStreamReader(self, schema: StructType):
    #    return SimpleStreamReader()

    def streamWriter(self, schema: StructType, overwrite: bool):
        return FakeStreamWriter(self.options)

Krok 2: Implementace čtečky streamu

Dále implementujte příklad streamované čtečky dat, která vygeneruje dva řádky v každém mikrobatchu. Můžete implementovat DataSourceStreamReadernebo pokud má zdroj dat nízkou propustnost a nevyžaduje dělení, můžete místo toho implementovat SimpleDataSourceStreamReader . Buď simpleStreamReader() nebo streamReader() musí být implementována, a simpleStreamReader() je vyvolána pouze v případě, že streamReader() není implementována.

Implementace DataSourceStreamReader

Instance streamReader má celočíselné posuny, které se zvýší o 2 v každém mikrobatchu implementovaném s rozhraním DataSourceStreamReader .

from pyspark.sql.datasource import InputPartition
from typing import Iterator, Tuple
import os
import json

class RangePartition(InputPartition):
    def __init__(self, start, end):
        self.start = start
        self.end = end

class FakeStreamReader(DataSourceStreamReader):
    def __init__(self, schema, options):
        self.current = 0

    def initialOffset(self) -> dict:
        """
        Returns the initial start offset of the reader.
        """
        return {"offset": 0}

    def latestOffset(self) -> dict:
        """
        Returns the current latest offset that the next microbatch will read to.
        """
        self.current += 2
        return {"offset": self.current}

    def partitions(self, start: dict, end: dict):
        """
        Plans the partitioning of the current microbatch defined by start and end offset. It
        needs to return a sequence of :class:`InputPartition` objects.
        """
        return [RangePartition(start["offset"], end["offset"])]

    def commit(self, end: dict):
        """
        This is invoked when the query has finished processing data before end offset. This
        can be used to clean up the resource.
        """
        pass

    def read(self, partition) -> Iterator[Tuple]:
        """
        Takes a partition as an input and reads an iterator of tuples from the data source.
        """
        start, end = partition.start, partition.end
        for i in range(start, end):
            yield (i, str(i))

Implementace SimpleDataSourceStreamReader

Instance SimpleStreamReader je stejná jako FakeStreamReader instance, která generuje dva řádky v každé dávce, ale implementovaná s rozhraním SimpleDataSourceStreamReader bez dělení.

class SimpleStreamReader(SimpleDataSourceStreamReader):
    def initialOffset(self):
        """
        Returns the initial start offset of the reader.
        """
        return {"offset": 0}

    def read(self, start: dict) -> (Iterator[Tuple], dict):
        """
        Takes start offset as an input, then returns an iterator of tuples and the start offset of the next read.
        """
        start_idx = start["offset"]
        it = iter([(i,) for i in range(start_idx, start_idx + 2)])
        return (it, {"offset": start_idx + 2})

    def readBetweenOffsets(self, start: dict, end: dict) -> Iterator[Tuple]:
        """
        Takes start and end offset as inputs, then reads an iterator of data deterministically.
        This is called when the query replays batches during restart or after a failure.
        """
        start_idx = start["offset"]
        end_idx = end["offset"]
        return iter([(i,) for i in range(start_idx, end_idx)])

    def commit(self, end):
        """
        This is invoked when the query has finished processing data before end offset. This can be used to clean up resources.
        """
        pass

Krok 3: Implementace zapisovače streamu

Teď implementujte zapisovač streamování. Tento zapisovač streamovaných dat zapisuje informace o metadatech jednotlivých mikrobatchů do místní cesty.

class SimpleCommitMessage:
   def __init__(self, partition_id: int, count: int):
       self.partition_id = partition_id
       self.count = count

class FakeStreamWriter(DataSourceStreamWriter):
   def __init__(self, options):
       self.options = options
       self.path = self.options.get("path")
       assert self.path is not None

   def write(self, iterator):
       """
       Writes the data, then returns the commit message of that partition. Library imports must be within the method.
       """
       from pyspark import TaskContext
       context = TaskContext.get()
       partition_id = context.partitionId()
       cnt = 0
       for row in iterator:
           cnt += 1
       return SimpleCommitMessage(partition_id=partition_id, count=cnt)

   def commit(self, messages, batchId) -> None:
       """
       Receives a sequence of :class:`WriterCommitMessage` when all write tasks have succeeded, then decides what to do with it.
       In this FakeStreamWriter, the metadata of the microbatch(number of rows and partitions) is written into a JSON file inside commit().
       """
       status = dict(num_partitions=len(messages), rows=sum(m.count for m in messages))
       with open(os.path.join(self.path, f"{batchId}.json"), "a") as file:
           file.write(json.dumps(status) + "\n")

   def abort(self, messages, batchId) -> None:
       """
       Receives a sequence of :class:`WriterCommitMessage` from successful tasks when some other tasks have failed, then decides what to do with it.
       In this FakeStreamWriter, a failure message is written into a text file inside abort().
       """
       with open(os.path.join(self.path, f"{batchId}.txt"), "w") as file:
           file.write(f"failed in batch {batchId}")

Krok 4: Registrace a použití ukázkového zdroje dat

Pokud chcete zdroj dat použít, zaregistrujte ho. Jakmile je regsitered, můžete ho použít v dotazech streamování jako zdroj nebo jímka předáním krátkého názvu nebo úplného názvu format(). Následující příklad zaregistruje zdroj dat a pak spustí dotaz, který čte z ukázkového zdroje dat a výstupy do konzoly:

spark.dataSource.register(FakeStreamDataSource)
query = spark.readStream.format("fakestream").load().writeStream.format("console").start()

Případně následující příklad používá ukázkový datový proud jako jímku a určuje výstupní cestu:

query = spark.readStream.format("fakestream").load().writeStream.format("fake").start("/output_path")

Řešení problému

Pokud je výstup následující chybou, výpočetní prostředky nepodporují vlastní zdroje dat PySpark. Musíte použít Databricks Runtime 15.2 nebo novější.

Error: [UNSUPPORTED_FEATURE.PYTHON_DATA_SOURCE] The feature is not supported: Python data sources. SQLSTATE: 0A000