Sources de données personnalisées PySpark

Les sources de données personnalisées PySpark sont créées en utilisant l'API DataSource Python (PySpark), ce qui permet de lire à partir de sources de données personnalisées et d'écrire dans des destinations de données personnalisées dans Apache Spark en utilisant Python. Vous pouvez utiliser les sources de données personnalisées PySpark pour définir des connexions personnalisées aux systèmes de données et implémenter des fonctionnalités supplémentaires afin de créer des sources de données réutilisables.

remarque

Les sources de données personnalisées PySpark nécessitent Databricks Runtime 15.4 LTS et versions ultérieures, ou version 2 de l'environnement Serverless.

Classe DataSource

La DataSource PySpark est une classe de base qui fournit des méthodes pour créer des lecteurs et des rédacteurs de données.

Implémenter la sous-classe de la source de données

Selon votre cas d'utilisation, les éléments suivants doivent être implémentés par toute sous-classe pour rendre une source de données lisible, inscriptible, ou les deux :

Propriété ou Méthode	Description
`name`	Obligatoire. Le nom de la source de données
`schema`	Obligatoire. Le schéma de la source de données à lire ou à écrire
`reader()`	Doit retourner un `DataSourceReader` pour rendre la source de données lisible (batch)
`writer()`	Doit renvoyer un `DataSourceWriter` pour que le récepteur de données soit accessible en écriture (batch).
`streamReader()` OU `simpleStreamReader()`	Doit renvoyer un `DataSourceStreamReader` pour rendre le flux de données lisible (streaming)
`streamWriter()`	Doit renvoyer un `DataSourceStreamWriter` pour rendre le data stream inscriptible (streaming)

Propriété ou Méthode	Description
`name`	Obligatoire. Le nom de la source de données
`schema`	Obligatoire. Le schéma de la source de données à lire ou à écrire
`reader()`	Doit retourner un `DataSourceReader` pour rendre la source de données lisible (batch)
`writer()`	Doit renvoyer un `DataSourceWriter` pour que le récepteur de données soit accessible en écriture (batch).
`streamReader()` OU `simpleStreamReader()`	Doit renvoyer un `DataSourceStreamReader` pour rendre le flux de données lisible (streaming)
`streamWriter()`	Doit renvoyer un `DataSourceStreamWriter` pour rendre le data stream inscriptible (streaming)

remarque

Les DataSource, DataSourceReader, DataSourceWriter, DataSourceStreamReader, DataSourceStreamWriter définis par l'utilisateur et leurs méthodes doivent être sérialisables. En d'autres termes, ils doivent être un dictionnaire ou un dictionnaire imbriqué qui contient un type primitif.

Enregistrer la source de données

Après avoir implémenté l'interface, vous devez l'enregistrer, puis vous pouvez la charger ou l'utiliser d'une autre manière, comme indiqué dans l'exemple suivant :

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

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

Exemple 1 : Créer une source de données PySpark pour les requêtes batch

Pour démontrer les capacités de lecture de source de données PySpark, créez une source de données qui génère des données d'exemple à l'aide du package Python faker. Pour plus d'information sur faker, consultez la documentation Faker.

Installez le package faker à l'aide de la commande suivante :

Python
%pip install faker

Étape 1 : Implémenter le lecteur pour une query par batch

Tout d'abord, implémentez la logique du lecteur pour générer des exemples de données. Utilisez la bibliothèque faker installée pour remplir chaque champ du schéma.

Python
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)

Étape 2 : définir la source de données d'exemple

Ensuite, définissez votre nouvelle source de données PySpark en tant que sous-classe de DataSource avec un nom, un schéma et un lecteur. La méthode reader() doit être définie pour lire à partir d'une source de données dans une query batch.

Python
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)

Étape 3 : Enregistrer et utiliser l'exemple de source de données

Pour utiliser la source de données, enregistrez-la. By default, le FakeDataSource comporte trois lignes, et le schéma inclut ces string champs : name, date, zipcode, state. L'exemple suivant enregistre, charge et affiche la source de données d'exemple avec les default :

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

Output
+-----------------+----------+-------+----------+
|             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|
+-----------------+----------+-------+----------+

Seuls les champs string sont pris en charge, mais vous pouvez spécifier un schéma avec n'importe quels champs correspondant aux champs des fournisseurs de packages faker pour générer des données aléatoires à des fins de test et de développement. L'exemple suivant charge la source de données avec les champs name et company :

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

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

Pour charger la source de données avec un nombre personnalisé de lignes, spécifiez l'option numRows. L’exemple suivant spécifie 5 lignes :

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

Output
+--------------+----------+-------+------------+
|          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|
+--------------+----------+-------+------------+

Exemple 2 : Créer une source de données GitHub PySpark à l'aide de variantes

Pour démontrer l'utilisation des variantes dans une source de données PySpark, cet exemple crée une source de données qui lit les demandes d'extraction depuis GitHub.

remarque

Les variantes sont prises en charge avec les sources de données personnalisées PySpark dans Databricks Runtime 17.1 et versions supérieures.

Pour des informations sur les variantes, consultez query les données de variante.

Étape 1 : Mettez en œuvre le lecteur pour récupérer les requêtes pull

Tout d’abord, implémentez la logique de lecture pour récupérer les pull requests à partir du repository GitHub spécifié.

Python
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")
            )

Étape 2 : Définir le DataSource GitHub

Ensuite, définissez votre nouvelle source de données PySpark GitHub comme une sous-classe de DataSource avec un nom, un schéma et une méthode reader(). Le schéma inclut ces champs : id, title, user, created_at, updated_at. Le champ user est défini comme une variante.

Python
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)

Étape 3 : Enregistrer et utiliser la source de données

Pour utiliser la source de données, enregistrez-la. L'exemple suivant enregistre, puis charge la source de données et génère trois lignes des données de PR du repository GitHub :

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

Output
+---------+-----------------------------------------------------+---------------------+----------------------+----------------------+
| 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 |
+---------+-----------------------------------------------------+---------------------+----------------------+----------------------+

Exemple 3 : Créer une source de données PySpark pour la lecture et l'écriture en streaming

Pour démontrer les capacités de lecture et d'écriture de Stream de PySpark DataSource, créez une source de données d'exemple qui génère deux lignes dans chaque micro-batch à l'aide du package Python faker. Pour plus d'information sur faker, consultez la documentation Faker.

Installez le package faker à l'aide de la commande suivante :

Python
%pip install faker

Étape 1 : Implémenter le lecteur de Stream

Tout d'abord, implémentez l'exemple de lecteur de données en streaming qui génère deux lignes dans chaque microbatch. Vous pouvez implémenter DataSourceStreamReader, ou si la source de données a un faible throughput et ne nécessite pas de partitionnement, vous pouvez implémenter SimpleDataSourceStreamReader à la place. simpleStreamReader() ou streamReader() doivent être implémentés, et simpleStreamReader() n'est invoqué que lorsque streamReader() n'est pas implémenté.

Implémentation de DataSourceStreamReader

L'instance streamReader possède un décalage entier qui augmente de 2 dans chaque micro-lot, implémenté avec l'interface DataSourceStreamReader.

Python
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))

Implémentation de SimpleDataSourceStreamReader

L'instance SimpleStreamReader est la même que l'instance FakeStreamReader qui génère deux lignes dans chaque batch, mais implémentée avec l'interface SimpleDataSourceStreamReader sans partitionnement.

Python
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

Étape 2 : Implémenter l'enregistreur de Stream

Ensuite, implémentez le writer de streaming. Cet écrivain de données en streaming écrit les informations de métadonnées de chaque microbatch vers un chemin local.

Python
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}")

Étape 3 : Définir la Source de données d’exemple

Définissez maintenant votre nouvelle DataSource PySpark comme une sous-classe de DataSource avec un nom, un schéma et les méthodes streamReader() et streamWriter().

Python
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)

Étape 4 : enregistrez et utilisez la source de données exemple

Pour utiliser la source de données, enregistrez-la. Une fois enregistré, vous pouvez l'utiliser dans des queries streaming comme source ou sink en passant un nom court ou un nom complet à format(). L'exemple suivant enregistre la source de données, puis start une query qui lit à partir de la source de données exemple et sort vers la console :

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

Vous pouvez également utiliser le code suivant avec le stream d’exemple comme sink, et il spécifie un chemin de sortie :

Python
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()
)

Exemple 4 : Créez un connecteur de streaming Google BigQuery

L'exemple suivant montre comment créer un connecteur de streaming personnalisé pour Google BigQuery (BQ) à l'aide d'une source de données PySpark. Databricks fournit un connecteur Spark pour l'ingestion par batch BigQuery, et la Lakehouse Federation peut également se connecter à distance à n'importe quel ensemble de données BigQuery et extraire des données par la création de catalogues étrangers, mais aucun ne prend entièrement en charge les workflows de streaming incrémentiels ou continus. Ce connecteur permet une migration incrémentielle des données par phases et une migration quasi en temps réel à partir des tables BigQuery alimentées par des sources de streaming avec point de contrôle persistant.

Ce connecteur personnalisé a les fonctionnalités suivantes :

Compatible avec Structured Streaming et Lakeflow Pipelines.
Prend en charge le suivi incrémentiel des enregistrements et l'ingestion en streaming continue, et respecte la sémantique de Structured Streaming.
Utilise l'API BigQuery Storage avec un protocole basé sur RPC pour une transmission des données plus rapide et moins chère.
Écrit les tables migrées directement dans Unity Catalog.
Gère automatiquement les points de contrôle à l'aide d'un champ incrémentiel basé sur une date ou un timestamp.
Prend en charge l'ingestion par batch avec Trigger.AvailableNow().
Ne nécessite aucun stockage cloud intermédiaire.
Sérialise la transmission de données BigQuery au format Arrow ou Avro.
Gère l'auto-parallélisme et distribue le travail entre les Workers Spark en fonction du volume de données.
Convient pour la migration des couches Raw et Bronze depuis BigQuery, avec prise en charge de la migration des couches Silver et Gold à l'aide de modèles SCD de type 1 ou de type 2.

Prérequis

Avant de mettre en œuvre le connecteur personnalisé, installez les packages requis :

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

Étape 1 : Implémenter le lecteur de Stream

Premièrement, implémentez le lecteur de données de streaming. La sous-classe DataSourceStreamReader doit implémenter les méthodes suivantes :

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

Pour plus de détails sur chaque méthode, consultez Méthodes.

Python
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={
                &quot;parent&quot;: parent,
                &quot;read_session&quot;: requested_session,
                &quot;max_stream_count&quot;: 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

Étape 2 : Définissez la source de données

Ensuite, définissez la source de données personnalisée. La sous-classe DataSource doit implémenter les méthodes suivantes :

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

Pour plus de détails sur chaque méthode, consultez Méthodes.

Python
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)

Étape 3 : Configurez et start la query streaming

Enfin, enregistrez le connecteur, puis configurez et start la query en streaming :

Python
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()
)

Ordre d’exécution

L'ordre d'exécution de la fonction du Stream personnalisé est décrit ci-dessous.

Pour charger le DataFrame Spark Stream :

Python
name(cls)
schema()

Pour les micro-lots (n) d'un nouveau start de query ou lors du redémarrage d'une query existante (point de contrôle nouveau ou existant) :

Python
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()

Pour le prochain micro-lot (n+1) d'une query en cours d'exécution sur un point de contrôle existant :

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

remarque

La fonction latestOffset orchestre le checkpointing. Partagez une variable de point de contrôle de type primitif entre les fonctions et renvoyez-la comme un dictionnaire. Par exemple : return {"offset": str(self.last_offset)}

Exemple 5 : s'authentifier avec une API externe

Cet exemple montre comment authentifier une source de données PySpark avec une API HTTP externe à l'aide d'une connexion HTTP Unity Catalog, afin que le code de la source de données ne contienne jamais de jetons ou d'informations d'identification codés en dur.

remarque

L'injection d'informations d'identification de connexion HTTP d'Unity Catalog nécessite Databricks Runtime 18.1 ou une version ultérieure.

Étape 1 : Créer une connexion HTTP

Avant d'implémenter la source de données, créez une connexion HTTP nommée my_weather_api dans Unity Catalog et accordez aux utilisateurs ou aux groupes la permission MANAGE de l'utiliser. Seuls les utilisateurs disposant de l'autorisation MANAGE sur la connexion peuvent Trigger l'injection d'informations d'identification.

Stockez le jeton d'API en tant que secret Databricks et référencez-le avec la fonction secret plutôt que de saisir le jeton littéral, afin que l'identifiant n'apparaisse jamais dans la définition de connexion.

SQL
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`;

Étape 2 : Implémentez le lecteur pour une query batch

Ensuite, implémentez la logique du lecteur pour récupérer les données de l'API REST. Le lecteur lit les valeurs injectées host, base_path et bearer_token à partir de ses options, ainsi aucun identifiant n'apparaît dans le code.

Python
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"])

Étape 3 : Définir la Source de données d’exemple

Définissez maintenant votre nouvelle Source de données PySpark en tant que sous-classe de DataSource avec un nom, un schéma et un lecteur.

Python
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)

Étape 4 : Enregistrer et utiliser la source de données

Pour utiliser la source de données, enregistrez-la. Référencez ensuite la connexion HTTP Unity Catalog avec l'option databricks.connection. Le driver Spark récupère automatiquement les informations d'identification OAuth2 de courte durée depuis Unity Catalog et les injecte (par exemple, bearer_token, host et base_path) dans la carte des options de la source de données. Les clés d'identification injectées par Unity Catalog ne peuvent pas être remplacées, et les options qui sont globalement bloquées, telles que host et port, restent bloquées et ne peuvent pas être définies par les utilisateurs.

Python
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()

Cet exemple implémente uniquement des lectures par batch. La même option databricks.connection s'applique également aux lectures et écritures en streaming lorsque votre source de données implémente les méthodes correspondantes (streamReader ou simpleStreamReader pour les lectures en streaming, et writer ou streamWriter pour les écritures).

Dépannage

Si la sortie est l’erreur suivante, votre compute ne prend pas en charge les sources de données personnalisées PySpark. Vous devez utiliser Databricks Runtime 15.2 ou une version ultérieure.

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

Classe DataSource​

Implémenter la sous-classe de la source de données​

Enregistrer la source de données​

Exemple 1 : Créer une source de données PySpark pour les requêtes batch​

Étape 1 : Implémenter le lecteur pour une query par batch​

Étape 2 : définir la source de données d'exemple​

Étape 3 : Enregistrer et utiliser l'exemple de source de données​

Exemple 2 : Créer une source de données GitHub PySpark à l'aide de variantes​

Étape 1 : Mettez en œuvre le lecteur pour récupérer les requêtes pull​

Étape 2 : Définir le DataSource GitHub​

Étape 3 : Enregistrer et utiliser la source de données​

Exemple 3 : Créer une source de données PySpark pour la lecture et l'écriture en streaming​

Étape 1 : Implémenter le lecteur de Stream​

Implémentation de DataSourceStreamReader​

Implémentation de SimpleDataSourceStreamReader​

Étape 2 : Implémenter l'enregistreur de Stream​

Étape 3 : Définir la Source de données d’exemple​

Étape 4 : enregistrez et utilisez la source de données exemple​

Exemple 4 : Créez un connecteur de streaming Google BigQuery​

Prérequis​

Étape 1 : Implémenter le lecteur de Stream​

Étape 2 : Définissez la source de données​

Étape 3 : Configurez et start la query streaming​

Ordre d’exécution​

Exemple 5 : s'authentifier avec une API externe​

Étape 1 : Créer une connexion HTTP​

Étape 2 : Implémentez le lecteur pour une query batch​

Étape 3 : Définir la Source de données d’exemple​

Étape 4 : Enregistrer et utiliser la source de données​

Dépannage​

Classe DataSource

Implémenter la sous-classe de la source de données

Enregistrer la source de données

Exemple 1 : Créer une source de données PySpark pour les requêtes batch

Étape 1 : Implémenter le lecteur pour une query par batch

Étape 2 : définir la source de données d'exemple

Étape 3 : Enregistrer et utiliser l'exemple de source de données

Exemple 2 : Créer une source de données GitHub PySpark à l'aide de variantes

Étape 1 : Mettez en œuvre le lecteur pour récupérer les requêtes pull

Étape 2 : Définir le DataSource GitHub

Étape 3 : Enregistrer et utiliser la source de données

Exemple 3 : Créer une source de données PySpark pour la lecture et l'écriture en streaming

Étape 1 : Implémenter le lecteur de Stream

Implémentation de DataSourceStreamReader

Implémentation de SimpleDataSourceStreamReader

Étape 2 : Implémenter l'enregistreur de Stream

Étape 3 : Définir la Source de données d’exemple

Étape 4 : enregistrez et utilisez la source de données exemple

Exemple 4 : Créez un connecteur de streaming Google BigQuery

Prérequis

Étape 1 : Implémenter le lecteur de Stream

Étape 2 : Définissez la source de données

Étape 3 : Configurez et start la query streaming

Ordre d’exécution

Exemple 5 : s'authentifier avec une API externe

Étape 1 : Créer une connexion HTTP

Étape 2 : Implémentez le lecteur pour une query batch

Étape 3 : Définir la Source de données d’exemple

Étape 4 : Enregistrer et utiliser la source de données

Dépannage