Vlastní zdroje dat PySpark

Vlastní zdroje dat PySpark se vytvářejí pomocí rozhraní API Zdroje dat Pythonu (PySpark), 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í pro vytváření opakovaně použitelných zdrojů dat.

Poznámka:

Vlastní zdroje dat PySpark vyžadují Databricks Runtime 15.4 LTS a vyšší nebo bezserverové prostředí verze 2.

třída DataSource

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 konkrétním případě použití musí každá podtřída implementovat následující, aby byl zdroj dat buď č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ávkový).
writer()	Musí se vrátit DataSourceWriter, aby byla datová sink zapisovatelná (dávkově).
streamReader() nebo simpleStreamReader()	Je třeba vrátit DataSourceStreamReader, aby byl datový proud čitelný (streamování).
streamWriter()	Musí se vrátit DataSourceStreamWriter, aby byl datový proud zapisovatelný (streamování).

Poznámka:

Uživatelem definované DataSource, DataSourceReader, DataSourceWriter, DataSourceStreamReader, DataSourceStreamWriter a jejich metody musí být serializovány. 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: Implementujte čtečku pro dávkový dotaz

Nejprve implementujte logiku čtenáře pro generování ukázkových dat. 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 2: Definování ukázkového zdroje dat

Dále definujte nový zdroj dat PySpark jako podtřídu DataSource s názvem, schématem a čtenářem. Metoda reader() musí být definována 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 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 a company poli

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 GitHubu PySpark pomocí variant

Abychom si ukázali použití variant v datovém zdroji PySpark, vytvoří tento příklad zdroj dat, který čte žádosti o přijetí změn z GitHubu.

Poznámka:

Varianty jsou podporovány u vlastních zdrojů dat PySpark v Databricks Runtime 17.1 a novějších.

Informace o variantách najdete v tématu Dotazování na data variant.

Krok 1: Implementace čtečky pro načtení žádostí o přijetí změn

Nejprve implementujte logiku čtenáře pro načtení žádostí o přijetí změn ze zadaného úložiště GitHub.

class GithubVariantPullRequestReader(DataSourceReader):
    def __init__(self, options):
        self.token = options.get("token")
        self.repo = options.get("path")
        if self.repo is None:
            raise Exception(f"Must specify a repo in `.load()` method.")
        # Every value in this `self.options` dictionary is a string.
        self.num_rows = int(options.get("numRows", 10))

    def read(self, partition):
        header = {
            "Accept": "application/vnd.github+json",
        }
        if self.token is not None:
            header["Authorization"] = f"Bearer {self.token}"
        url = f"https://api.github.com/repos/{self.repo}/pulls"
        response = requests.get(url, headers=header)
        response.raise_for_status()
        prs = response.json()
        for pr in prs[:self.num_rows]:
            yield Row(
                id = pr.get("number"),
                title = pr.get("title"),
                user = VariantVal.parseJson(json.dumps(pr.get("user"))),
                created_at = pr.get("created_at"),
                updated_at = pr.get("updated_at")
            )

Krok 2: Definování Zdroje dat GitHubu

Dále definujte nový zdroj dat PySpark GitHub jako podtřídu DataSource s názvem, schématem a metodou reader(). Schéma obsahuje tato pole: id, titleuser, created_at, , . updated_at Pole user je definováno jako varianta.

import json
import requests

from pyspark.sql import Row
from pyspark.sql.datasource import DataSource, DataSourceReader
from pyspark.sql.types import VariantVal

class GithubVariantDataSource(DataSource):
    @classmethod
    def name(self):
        return "githubVariant"
    def schema(self):
        return "id int, title string, user variant, created_at string, updated_at string"
    def reader(self, schema):
        return GithubVariantPullRequestReader(self.options)

Krok 3: Registrace a použití zdroje dat

Pokud chcete zdroj dat použít, zaregistrujte ho. Následující příklad nejprve zaregistruje, poté načte zdroj dat a vypíše tři řádky dat pull requestů z úložiště na GitHubu.

spark.dataSource.register(GithubVariantDataSource)
spark.read.format("githubVariant").option("numRows", 3).load("apache/spark").display()

+---------+-----------------------------------------------------+---------------------+----------------------+----------------------+
| id      | title                                               | user                | created_at           | updated_at           |
+---------+---------------------------------------------------- +---------------------+----------------------+----------------------+
|   51293 |[SPARK-52586][SQL] Introduce AnyTimeType             |  {"avatar_url":...} | 2025-06-26T09:20:59Z | 2025-06-26T15:22:39Z |
|   51292 |[WIP][PYTHON] Arrow UDF for aggregation              |  {"avatar_url":...} | 2025-06-26T07:52:27Z | 2025-06-26T07:52:37Z |
|   51290 |[SPARK-50686][SQL] Hash to sort aggregation fallback |  {"avatar_url":...} | 2025-06-26T06:19:58Z | 2025-06-26T06:20:07Z |
+---------+-----------------------------------------------------+---------------------+----------------------+----------------------+

Příklad 3: Vytvoření DataSource PySpark pro streamované čtení a zápis

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: Implementace čtečky streamu

Nejprve implementujte příklad streamované čtečky dat, která generuje 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ý posun, který se zvýší o 2 v každém mikrobatchi, 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

Instanci SimpleStreamReader lze považovat za stejnou jako instanci FakeStreamReader, která generuje dva řádky v každé dávce, ale je implementována pomocí rozhraní 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 2: Implementace zapisovače streamu

Dále implementujte zapisovač pro streamování. Tento zapisovač streamovaných dat zaznamenává informace o metadatech každé mikro dávky do lokálního adresáře.

from pyspark.sql.datasource import DataSourceStreamWriter, WriterCommitMessage

class SimpleCommitMessage(WriterCommitMessage):
   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 and then returns the commit message for 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 3: Definování ukázkového zdroje dat

Nyní definujte nový PySpark DataSource 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 4: Registrace a použití ukázkového zdroje dat

Pokud chcete zdroj dat použít, zaregistrujte ho. Jakmile je zaregistrován, můžete ho použít v dotazech streamování buď jako zdroj, nebo jako výstup tím, že předáte krátký nebo plný název do 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í kód používá ukázkový stream jako jímku a určuje výstupní cestu:

spark.dataSource.register(FakeStreamDataSource)

# Make sure the output directory exists and is writable
output_path = "/output_path"
dbutils.fs.mkdirs(output_path)
checkpoint_path = "/output_path/checkpoint"

query = (
    spark.readStream
    .format("fakestream")
    .load()
    .writeStream
    .format("fakestream")
    .option("path", output_path)
    .option("checkpointLocation", checkpoint_path)
    .start()
)

Příklad 4: Vytvoření streamovacího konektoru Google BigQuery

Následující příklad ukazuje, jak vytvořit vlastní streamovací konektor pro Google BigQuery (BQ) pomocí PySpark DataSource. Databricks poskytuje konektor Sparku pro dávkový příjem dat BigQuery a Lakehouse Federation se také může vzdáleně připojit k jakékoli datové sadě BigQuery a načíst data prostřednictvím vytvoření cizího katalogu, ale ani jeden z nich plně nepodporuje přírůstkové nebo průběžné streamovací pracovní postupy. Tento konektor umožňuje přírůstkovou migraci fázových dat a migraci z tabulek BigQuery v téměř reálném čase živených zdroji streamování s trvalým kontrolním bodem.

Tento vlastní konektor má následující funkce:

• Kompatibilní se Structured Streaming a pipelinemi Lakeflow.
• Podporuje sledování přírůstkových záznamů a průběžný příjem streamů a řídí se sémantikou strukturovaného streamování.
• K rychlejšímu a levnějšímu přenosu dat používá rozhraní API služby BigQuery Storage s protokolem ZALOŽENÝm na RPC.
• Zapisuje migrované tabulky přímo do katalogu Unity.
• Spravuje kontrolní body automaticky pomocí přírůstkového pole založeného na datu nebo časovém razítku.
• Podporuje dávkové zpracování s Trigger.AvailableNow().
• Nevyžaduje žádné zprostředkující cloudové úložiště.
• Serializuje přenos dat BigQuery pomocí formátu Arrow nebo Avro.
• Zpracovává automatické paralelizace a distribuuje práci mezi pracovní uzly Spark na základě objemu dat.
• Vhodné pro migraci surové a bronzové vrstvy z BigQuery s podporou migrace vrstev Silver a Gold pomocí vzorů SCD Type 1 nebo Type 2.

Předpoklady

Před implementací vlastního konektoru nainstalujte požadované balíčky:

%pip install faker google.cloud google.cloud.bigquery google.cloud.bigquery_storage

Krok 1: Implementace čtečky streamu

Nejprve implementujte čtečku streamovaných dat. Podtřída DataSourceStreamReader musí implementovat následující metody:

• initialOffset(self) -> dict
• latestOffset(self) -> dict
• partitions(self, start: dict, end: dict) -> Sequence[InputPartition]
• read(self, partition: InputPartition) -> Union[Iterator[Tuple], Iterator[Row]]
• commit(self, end: dict) -> None
• stop(self) -> None

Podrobnosti o jednotlivých metodách naleznete v tématu Metody.

import os
from pyspark.sql.datasource import DataSourceStreamReader, InputPartition
from pyspark.sql.datasource import DataSourceStreamWriter
from pyspark.sql import Row
from pyspark.sql import SparkSession
from pyspark.sql.datasource import DataSource
from pathlib import Path
from pyarrow.lib import TimestampScalar
from datetime import datetime
from typing import Iterator, Tuple, Any, Dict, List, Sequence
from google.cloud.bigquery_storage import BigQueryReadClient, ReadSession
from google.cloud import bigquery_storage
import pandas
import datetime
import uuid
import time, logging

start_time = time.time()


class RangePartition(InputPartition):
    def __init__(self, session: ReadSession, stream_idx: int):
        self.session = session
        self.stream_idx = stream_idx


class BQStreamReader(DataSourceStreamReader):

    def __init__(self, schema, options):
        self.project_id = options.get("project_id")
        self.dataset = options.get("dataset")
        self.table = options.get("table")
        self.json_auth_file = "/home/"+options.get("service_auth_json_file_name")
        self.max_parallel_conn = options.get("max_parallel_conn", 1000)
        self.incremental_checkpoint_field = options.get("incremental_checkpoint_field", "")

        self.last_offset = None

    def initialOffset(self) -> dict:
        """
        Returns the initial start offset of the reader.
        """
        from datetime import datetime
        logging.info("Inside initialOffset!!!!!")
        # self.increment_latest_vals.append(datetime.strptime('1900-01-01 23:57:12', "%Y-%m-%d %H:%M:%S"))
        self.last_offset = '1900-01-01 23:57:12'

        return {"offset": str(self.last_offset)}

    def latestOffset(self):
        """
        Returns the current latest offset that the next microbatch will read to.
        """
        from datetime import datetime
        from google.cloud import bigquery

        if (self.last_offset is None):
            self.last_offset = '1900-01-01 23:57:12'

        client = bigquery.Client.from_service_account_json(self.json_auth_file)
        # max_offset=start["offset"]
        logging.info(f"************************last_offset: {self.last_offset}***********************")
        f_sql_str = ''
        for x_str in self.incremental_checkpoint_field.strip().split(","):
            f_sql_str += f"{x_str}>'{self.last_offset}' or "
        f_sql_str = f_sql_str[:-3]
        job_query = client.query(
            f"select max({self.incremental_checkpoint_field}) from {self.project_id}.{self.dataset}.{self.table} where {f_sql_str}")
        for query in job_query.result():
            max_res = query[0]

        if (str(max_res).lower() != 'none'):
            return {"offset": str(max_res)}

        return {"offset": str(self.last_offset)}

    def partitions(self, start: dict, end: dict) -> Sequence[InputPartition]:

        """
        Plans the partitioning of the current microbatch defined by start and end offset. It
        needs to return a sequence of :class:`InputPartition` objects.
        """
        if (self.last_offset is None):
            self.last_offset = end['offset']

        os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = self.json_auth_file

        # project_id = self.auth_project_id

        client = BigQueryReadClient()

        # This example reads baby name data from the public datasets.
        table = "projects/{}/datasets/{}/tables/{}".format(
            self.project_id, self.dataset, self.table
        )
        requested_session = bigquery_storage.ReadSession()
        requested_session.table = table
        if (self.incremental_checkpoint_field != ''):
            start_offset = start["offset"]
            end_offset = end["offset"]
            f_sql_str = ''
            for x_str in self.incremental_checkpoint_field.strip().split(","):
                f_sql_str += f"({x_str}>'{start_offset}' and {x_str}<='{end_offset}') or "
            f_sql_str = f_sql_str[:-3]
            requested_session.read_options.row_restriction = f"{f_sql_str}"

        # This example leverages Apache Avro.
        requested_session.data_format = bigquery_storage.DataFormat.AVRO

        parent = "projects/{}".format(self.project_id)
        session = client.create_read_session(
            request={
                "parent": parent,
                "read_session": requested_session,
                "max_stream_count": int(self.max_parallel_conn),
            },
        )
        self.last_offset = end['offset']
        return [RangePartition(session, i) for i in range(len(session.streams))]

    def read(self, partition) -> Iterator[List]:
        """
        Takes a partition as an input and reads an iterator of tuples from the data source.
        """
        from datetime import datetime
        session = partition.session
        stream_idx = partition.stream_idx
        os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = self.json_auth_file
        client_1 = BigQueryReadClient()
        # requested_session.read_options.selected_fields = ["census_tract", "clearance_date", "clearance_status"]
        reader = client_1.read_rows(session.streams[stream_idx].name)
        reader_iter = []

        for message in reader.rows():
            reader_iter_in = []
            for k, v in message.items():
                reader_iter_in.append(v)
            # yield(reader_iter)
            reader_iter.append(reader_iter_in)
            # yield (message['hash'], message['size'], message['virtual_size'], message['version'])
        # self.increment_latest_vals.append(max_incr_val)
        return iter(reader_iter)

    def commit(self, end):

        """
        This is invoked when the query has finished processing data before end offset. This
        can be used to clean up the resource.
        """
        pass

Krok 2: Definování zdroje dat

Dále definujte vlastní zdroj dat. Podtřída DataSource musí implementovat následující metody:

• name(cls) -> str
• schema(self) -> Union[StructType, str]

Podrobnosti o jednotlivých metodách naleznete v tématu Metody.

from pyspark.sql.datasource import DataSource
from pyspark.sql.types import StructType
from google.cloud import bigquery

class BQStreamDataSource(DataSource):
    """
    An example data source for streaming data from a public API containing users' comments.
    """

    @classmethod
    def name(cls):
        return "bigquery-streaming"

    def schema(self):
        type_map = {'integer': 'long', 'float': 'double', 'record': 'string'}
        json_auth_file = "/home/" + self.options.get("service_auth_json_file_name")
        client = bigquery.Client.from_service_account_json(json_auth_file)
        table_ref = self.options.get("project_id") + '.' + self.options.get("dataset") + '.' + self.options.get("table")
        table = client.get_table(table_ref)
        original_schema = table.schema
        result = []
        for schema in original_schema:
            col_attr_name = schema.name
            if (schema.mode != 'REPEATED'):
                col_attr_type = type_map.get(schema.field_type.lower(), schema.field_type.lower())
            else:
                col_attr_type = f"array<{type_map.get(schema.field_type.lower(), schema.field_type.lower())}>"
            result.append(col_attr_name + " " + col_attr_type)

        return ",".join(result)
        # return "census_tract double,clearance_date string,clearance_status string"

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

Krok 3: Konfigurace a spuštění streamovacího dotazu

Nakonec zaregistrujte konektor a pak nakonfigurujte a spusťte dotaz streamování:

spark.dataSource.register(BQStreamDataSource)

# Ingests table data incrementally using the provided timestamp-based field.
# The latest value is checkpointed using offset semantics.
# Without the incremental input field, full table ingestion is performed.
# Service account JSON files must be available to every Spark executor worker
# in the /home folder using --files /home/<file_name>.json or an init script.

query = (
    spark.readStream.format("bigquery-streaming")
    .option("project_id", <bq_project_id>)
    .option("incremental_checkpoint_field", <table_incremental_ts_based_col>)
    .option("dataset", <bq_dataset_name>)
    .option("table", <bq_table_name>)
    .option("service_auth_json_file_name", <service_account_json_file_name>)
    .option("max_parallel_conn", <max_parallel_threads_to_pull_data>)  # defaults to max 1000
    .load()
)

(
    query.writeStream.trigger(processingTime="30 seconds")
    .option("checkpointLocation", "checkpoint_path")
    .foreachBatch(writeToTable)  # your target table write function
    .start()
)

Pořadí provádění

Pořadí provádění funkce uživatelského datového proudu je popsáno níže.

Při načítání datového rámce pro stream Spark:

name(cls)
schema()

Pro mikrodávku (n) při spuštění nového dotazu nebo restartování existujícího dotazu (nový či existující kontrolní bod):

partitions(end_offset, end_offset)  # loads the last saved offset from the checkpoint at query restart
latestOffset()
partitions(start_offset, end_offset)  # plans partitions and distributes to Python workers
read()  # user’s source read definition, runs on each Python worker
commit()

Pro další mikrobatch (n+1) spuštěného dotazu na existujícím kontrolním bodu:

latestOffset()
partitions(start_offset, end_offset)
read()
commit()

Poznámka:

Funkce latestOffset orchestruje vytváření kontrolních bodů. Sdílejte proměnnou kontrolního bodu primitivního typu mezi funkcemi a vraťte ji jako slovník. Příklad: return {"offset": str(self.last_offset)}

Příklad 5: Ověření pomocí externího rozhraní API

Tento příklad ukazuje, jak ověřit zdroj dat PySpark pomocí externího rozhraní HTTP API pomocí připojení HTTP katalogu Unity, aby kód zdroje dat nikdy neobsahuje pevně zakódované tokeny nebo přihlašovací údaje.

Poznámka:

Vkládání přihlašovacích údajů pro připojení HTTP v katalogu Unity vyžaduje Databricks Runtime verze 18.1 nebo novější.

Krok 1: Vytvoření připojení HTTP

Před implementací zdroje dat vytvořte připojení HTTP s názvem my_weather_api v katalogu Unity a udělte uživatelům nebo skupinám MANAGE oprávnění. Injektáž přihlašovacích údajů můžou aktivovat jenom uživatelé s MANAGE oprávněním k připojení.

Token rozhraní API uložte jako tajný údaj v Databricks a odkazujte na něj pomocí funkce secret namísto přímého zadání tokenu, aby se přihlašovací údaj nikdy nezobrazil v definici připojení.

CREATE CONNECTION my_weather_api TYPE HTTP
OPTIONS (
    host 'https://api.openweathermap.org',
    base_path '/data/2.5',
    bearer_token secret('my_secret_scope', 'weather_api_token')
);

GRANT MANAGE ON CONNECTION my_weather_api TO `user@example.com`;

Krok 2: Implementujte čtečku dávkového dotazu

Dále implementujte logiku čtenáře, která načítá data z rozhraní REST API. Čtečka načítá vložené hodnoty host, base_path a bearer_token ze svého nastavení, takže se v kódu neobjevují žádné přihlašovací údaje.

from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition
from urllib.parse import quote
import urllib.error
import urllib.request
import json

class WeatherApiReader(DataSourceReader):
    def __init__(self, options):
        self.host = options["host"]
        self.base_path = options["base_path"]
        self.token = options["bearer_token"]
        # Every value in this `options` dictionary is a string.
        self.cities = options.get("cities", "Seattle,Portland,Denver").split(",")

    def partitions(self):
        return [InputPartition(city.strip()) for city in self.cities]

    def read(self, partition):
        city = partition.value
        # URL-encode the city so names with spaces or non-ASCII characters (for example, "New York" or "São Paulo") produce a valid query string.
        url = f"{self.host}{self.base_path}/weather?q={quote(city)}&units=metric"
        req = urllib.request.Request(url)
        req.add_header("Authorization", f"Bearer {self.token}")
        try:
            # Set a timeout so a slow or unresponsive API surfaces a controlled error instead of hanging the Spark task.
            with urllib.request.urlopen(req, timeout=30) as resp:
                data = json.loads(resp.read().decode())
        except (urllib.error.URLError, TimeoutError) as e:
            raise RuntimeError(f"Weather API request failed for {city}: {e}")
        # Validate the response shape before indexing so an error payload raises a clear message instead of a KeyError.
        try:
            main = data["main"]
            weather = data["weather"][0]
        except (KeyError, IndexError, TypeError):
            raise RuntimeError(f"Unexpected weather API response for {city}: {data}")
        yield (city, main["temp"], main["humidity"], weather["description"])

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

Teď definujte nový zdroj dat PySpark jako podtřídu DataSource s názvem, schématem a čtenářem.

class WeatherApiSource(DataSource):
    def __init__(self, options):
        self.options = options

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

    def schema(self):
        return "city STRING, temperature DOUBLE, humidity INT, description STRING"

    def reader(self, schema):
        return WeatherApiReader(self.options)

Krok 4: Registrace a použití zdroje dat

Pokud chcete zdroj dat použít, zaregistrujte ho. Pak pomocí možnosti databricks.connection odkazujte na připojení HTTP ke službě Unity Catalog. Ovladač Sparku automaticky načte krátkodobé přihlašovací údaje OAuth2 z katalogu Unity a vloží je (například bearer_token, hosta base_path) do mapy možností zdroje dat. Klíče přihlašovacích údajů vložené do katalogu Unity nelze přepsat a možnosti, které jsou globálně blokované, například host a port, zůstanou blokované a uživatelé je nemůžou nastavit.

spark.dataSource.register(WeatherApiSource)

df = (
    spark.read.format("weather_api")
    .option("databricks.connection", "my_weather_api")   # Unity Catalog injects host, base_path, bearer_token
    .option("cities", "Seattle,Portland,Denver")         # user-defined option passes through
    .load()
)
df.show()

Tento příklad implementuje jen dávkové čtení. Stejná databricks.connection možnost platí také pro streamování čtení a zápisů, když zdroj dat implementuje odpovídající metody (streamReader nebo simpleStreamReader pro čtení streamování nebo writerstreamWriter pro zápisy).

Ř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