Fontes de dados personalizadas do PySpark

As fontes de dados personalizadas do PySpark são criadas usando a API DataSource do Python (PySpark), que permite a leitura de fontes de dados personalizadas e gravação em coletores de dados personalizados no Apache Spark usando Python. Você pode usar fontes de dados personalizadas do PySpark para definir conexões personalizadas com sistemas de dados e implementar funcionalidades adicionais para criar fontes de dados reutilizáveis.

O Spark inclui suporte embutido para formatos padrão como Delta, Iceberg, Parquet, JSON, CSV e JDBC, mas não para muitos outros sistemas como APIs REST, Google Sheets, conjuntos de dados Hugging Face ou serviços internos proprietários. A API Python DataSource preenche essa lacuna: você constrói conectores para esses sistemas em Python puro, sem desenvolvimento de conectores baseados em JVM, e os usa como qualquer fonte de dados Spark embutida, inclusive no Spark SQL.

Classe DataSource

O PySpark DataSource é uma classe base que fornece métodos para criar leitores e gravadores de dados.

Implementar a subclasse da fonte de dados

Dependendo do caso de uso, o seguinte deve ser implementado por qualquer subclasse para tornar uma fonte de dados legível, gravável ou ambas:

Registrar a fonte de dados

Após implementar a interface, você deverá registrá-la e, em seguida, pode carregá-la ou usá-la de outra forma, conforme mostrado no exemplo a seguir:

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

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

Exemplo 1: criar uma Fonte de dados do PySpark para consulta em lote

Para demonstrar os recursos de leitor da Fonte de dados do PySpark, crie uma fonte de dados que gere dados de exemplo usando o pacote do Python faker. Para obter mais informações sobre faker, consulte a documentação de Faker.

Instale o pacote faker usando o seguinte comando:

%pip install faker

Etapa 1: Implementar o leitor para uma consulta em lote

Primeiro, implemente a lógica do leitor para gerar dados de exemplo. Use a biblioteca faker instalada para preencher cada campo no esquema.

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)

Etapa 2: Definir o exemplo DataSource

Em seguida, defina seu novo PySpark DataSource como uma subclasse de DataSource com nome, esquema e leitor. O método reader() deve ser definido para ler de uma fonte de dados em uma consulta em lote.

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)

Etapa 3: registrar e usar a fonte de dados de exemplo

Para usar a fonte de dados, registre-a. Por padrão, o FakeDataSource tem três linhas, e o esquema inclui estes campos de string: name, date, zipcode, state. O exemplo a seguir registra, carrega e gera a fonte de dados de exemplo com os padrões:

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

Apenas campos string são suportados, mas você pode especificar um esquema com quaisquer campos que correspondam aos campos dos provedores de pacotes faker para gerar dados aleatórios para teste e desenvolvimento. O exemplo a seguir carrega a fonte de dados com campos name e company:

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

Para carregar a fonte de dados com um número personalizado de linhas, especifique a opção numRows. O exemplo a seguir especifica cinco linhas:

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

Exemplo 2: Gravar em um destino de dados personalizado em uma consulta em lote

Para demonstrar as capacidades de escrita da fonte de dados do PySpark, crie uma fonte de dados que grave cada partição de um DataFrame em um arquivo e, em seguida, grave um arquivo marcador de resumo quando o trabalho for confirmado.

Passo 1: Implemente o escritor para uma consulta em lote

Primeiro, implemente a lógica do escritor. Cada executor chama write() uma vez por partição. Após todas as operações de escrita serem concluídas com êxito, o driver chama commit(). Se alguma tarefa falhar, o motorista chama abort() em seu lugar.

from dataclasses import dataclass
from pyspark.sql.datasource import DataSourceWriter, WriterCommitMessage

@dataclass
class SimpleCommitMessage(WriterCommitMessage):
    partition_id: int
    count: int

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

    def write(self, iterator):
        """
        Writes the rows in a partition to a file, then returns a commit message with the row count. Library imports must be within the method.
        """
        import json
        import os
        from pyspark import TaskContext

        # Runs on an executor, so create the output directory on the local node.
        os.makedirs(self.path, exist_ok=True)
        partition_id = TaskContext.get().partitionId()
        count = 0
        with open(os.path.join(self.path, f"part-{partition_id}.json"), "w") as file:
            for row in iterator:
                file.write(json.dumps(row.asDict()) + "\n")
                count += 1
        return SimpleCommitMessage(partition_id=partition_id, count=count)

    def commit(self, messages):
        """
        Runs on the driver after all write tasks succeed. Writes a summary of the write to a marker file.
        """
        import json
        import os

        # Runs on the driver, so create the output directory on the local node.
        os.makedirs(self.path, exist_ok=True)
        total_rows = sum(message.count for message in messages if message is not None)
        with open(os.path.join(self.path, "_SUCCESS"), "w") as file:
            file.write(json.dumps({"partitions": len(messages), "rows": total_rows}))

    def abort(self, messages):
        """
        Runs on the driver if any write task fails. Use it to clean up partial output.
        """
        import os

        # Runs on the driver, so create the output directory on the local node.
        os.makedirs(self.path, exist_ok=True)
        with open(os.path.join(self.path, "_FAILED"), "w") as file:
            file.write("write job aborted")

Passo 2: Defina uma Fonte de Dados Gravável

Em seguida, defina uma DataSource subclasse que implemente writer(). O argumento overwrite é True quando o modo de escrita é overwrite e False quando é append.

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

class FakeSinkDataSource(DataSource):
    """
    An example writable data source that saves rows to files.
    """

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

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

    def writer(self, schema: StructType, overwrite: bool):
        return FakeDataSourceWriter(self.options)

Passo 3: Registrar e escrever no data sink

Para usar a fonte de dados, registre-a. Depois, escreva um DataFrame nele passando o nome curto para format() e um diretório de saída para a path opção. Este exemplo escreve em um caminho em um volume do Unity Catalog. Substitua <catalog>, <schema>, e <volume> por um volume existente.

spark.dataSource.register(FakeSinkDataSource)

output_path = "/Volumes/<catalog>/<schema>/<volume>/fakesink"

df = spark.range(3).selectExpr(
    "cast(id as string) as name",
    "'2025-01-01' as date",
    "'12345' as zipcode",
    "'California' as state",
)

df.write.format("fakesink").mode("append").option("path", output_path).save()

Exemplo 3: Crie uma fonte de dados PySpark no GitHub usando variantes

Para demonstrar o uso de variantes em um PySpark DataSource, este exemplo cria uma fonte de dados que lê solicitações de pull do GitHub.

Para obter informações sobre variantes, consulte Consulta de dados de variantes.

Etapa 1: Implementar o leitor para recuperar solicitações de pull

Primeiro, implemente a lógica do leitor para recuperar solicitações de pull do repositório GitHub especificado.

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

Etapa 2: Definir o GitHub DataSource

Em seguida, defina seu novo PySpark GitHub DataSource como uma subclasse de DataSource com nome, esquema e método reader(). O esquema inclui estes campos: id, , title, user, created_at, . updated_at O user campo é definido como uma variante.

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)

Etapa 3: Registrar e usar a fonte de dados

Para usar a fonte de dados, registre-a. O exemplo a seguir registra e carrega a fonte de dados e gera três linhas dos dados de PR do repositório GitHub:

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

Exemplo 4: Criar PySpark DataSource para leitura e gravação em streaming

Para demonstrar as capacidades de leitor e gravador de stream do PySpark DataSource, crie uma fonte de dados que gera duas linhas em cada microlote usando o pacote Python faker. Para obter mais informações sobre faker, consulte a documentação de Faker.

Instale o pacote faker usando o seguinte comando:

%pip install faker

Etapa 1: Implementar o leitor de fluxo

Primeiro, implemente o leitor de dados de streaming de exemplo que gera duas linhas em cada microbatch. Você pode implementar DataSourceStreamReaderou, se a fonte de dados tiver baixa taxa de transferência e não exigir particionamento, você poderá implementar SimpleDataSourceStreamReader . Tanto simpleStreamReader() como streamReader() deve ser implementado e simpleStreamReader() só é invocado quando streamReader() não é implementado.

Implementação do DataSourceStreamReader

A instância streamReader tem um deslocamento inteiro que aumenta duas vezes a cada microlote, implementado com a interface 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))

Implementação de SimpleDataSourceStreamReader

A instância SimpleStreamReader é a mesma que a instância FakeStreamReader que gera duas linhas em cada lote, mas implementada com a interface SimpleDataSourceStreamReader sem particionamento.

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

Etapa 2: Implementar o gravador de fluxo

Em seguida, implemente o gravador de streaming. Esse gravador de dados de streaming grava os metadados de cada microlote em um caminho local.

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

Etapa 3: Definir o exemplo DataSource

Agora, defina seu novo PySpark DataSource como uma subclasse de DataSource com um nome, um esquema e os métodos streamReader() e 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)

Etapa 4: registrar e usar a fonte de dados de exemplo

Para usar a fonte de dados, registre-a. Depois de registrado, ele pode ser utilizado em consultas de streaming como uma origem ou destino, passando um nome curto ou nome completo para format(). O exemplo a seguir registra a fonte de dados e, em seguida, inicia uma consulta que lê a fonte de dados de exemplo e envia a saída para o console.

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

Como alternativa, o código a seguir usa o fluxo de exemplo como um coletor e especifica um caminho de saída:

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

Exemplo 5: Crie um conector de streaming Google BigQuery

O exemplo a seguir demonstra como criar um conector de streaming personalizado para o Google BigQuery (BQ) usando um PySpark DataSource. O Databricks fornece um conector do Spark para ingestão em lotes do BigQuery, e a Lakehouse Federation também pode se conectar remotamente a qualquer conjunto de dados do BigQuery e extrair dados por meio da criação de catálogos externos, embora nenhum deles dê suporte a fluxos de trabalho de streaming incrementais ou contínuos. Esse conector permite a migração de dados incrementais em fases e a migração quase em tempo real de tabelas do BigQuery alimentadas por fontes de streaming com ponto de verificação persistente.

Esse conector personalizado tem os seguintes recursos:

• Compatível com pipelines de Streaming Estruturado e Lakeflow.
• Dá suporte ao acompanhamento de registros incrementais e à ingestão contínua de streaming e segue a semântica de Streaming Estruturado.
• Usa a API de Armazenamento do BigQuery com um protocolo baseado em RPC para transmissão de dados mais rápida e barata.
• Grava tabelas migradas diretamente para o Catálogo do Unity.
• Gerencia pontos de verificação automaticamente usando um campo incremental com base em data ou marcação de tempo.
• Dá suporte à ingestão em lote com Trigger.AvailableNow().
• Não requer armazenamento em nuvem intermediário.
• A serialização da transmissão de dados do BigQuery é feita usando o formato Arrow ou Avro.
• Lida com o paralelismo automático e distribui o trabalho entre os trabalhadores do Spark com base no volume de dados.
• Adequado para migração de camadas Raw e Bronze do BigQuery, com suporte para migrações de camadas Silver e Gold usando padrões SCD Tipo 1 ou Tipo 2.

Pré-requisitos

Antes de implementar o conector personalizado, instale os pacotes necessários:

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

Etapa 1: Implementar o leitor de fluxo

Primeiro, implemente o leitor de dados de streaming. A DataSourceStreamReader subclasse deve implementar os seguintes métodos:

• 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

Para obter detalhes sobre cada método, consulte Métodos.

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

Etapa 2: Definir o DataSource

Em seguida, defina a fonte de dados personalizada. A DataSource subclasse deve implementar os seguintes métodos:

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

Para obter detalhes sobre cada método, consulte Métodos.

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)

Etapa 3: Configurar e iniciar a consulta de streaming

Por fim, registre o conector e, em seguida, configure e inicie a consulta de streaming:

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

Ordem de execução

A ordem de execução da função do fluxo personalizado é descrita abaixo.

Para carregar o DataFrame do fluxo do Spark:

name(cls)
schema()

Para microbatch (n) de um novo início de consulta ou ao reiniciar uma consulta existente (ponto de verificação novo ou existente):

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

Para o próximo microlote (n+1) de uma consulta em execução em um ponto de verificação existente:

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

Exemplo 6: Autenticar com uma API externa

Este exemplo demonstra como autenticar uma fonte de dados do PySpark com uma API HTTP externa usando uma conexão HTTP do Catálogo do Unity, para que o código da fonte de dados nunca contenha tokens ou credenciais codificados.

Etapa 1: Criar uma conexão HTTP

Antes de implementar a fonte de dados, crie uma conexão HTTP nomeada my_weather_api no Catálogo do Unity e conceda permissão a usuários ou grupos MANAGE a ela. Somente usuários com a MANAGE permissão na conexão podem disparar a injeção de credencial.

Armazene o token de API como um segredo do Databricks e referencie-o com a secret função em vez de inserir o token literal, para que a credencial nunca apareça na definição de conexão.

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

Etapa 2: implementar o leitor para uma consulta em lote

Em seguida, implemente a lógica do leitor para buscar dados da API REST. O leitor lê os valores injetados host, base_path e bearer_token das suas opções, portanto nenhuma credencial aparece no código.

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

Etapa 3: Definir o exemplo DataSource

Agora, defina sua nova fonte de dados do PySpark como uma subclasse de DataSource, com um nome, um esquema e um leitor.

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)

Etapa 4: Registrar e usar a fonte de dados

Para usar a fonte de dados, registre-a. Em seguida, faça referência à conexão HTTP do Catálogo do Unity com a opção databricks.connection . O driver Spark recupera automaticamente as credenciais OAuth2 de curta duração do Catálogo do Unity e as injeta (por exemplo, bearer_token, hoste base_path) no mapa de opções da fonte de dados. As chaves de credencial injetadas pelo Catálogo do Unity não podem ser substituídas e as opções que estão bloqueadas globalmente, como host e port, permanecem bloqueadas e não podem ser definidas pelos usuários.

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

Este exemplo implementa somente leituras em lote. A mesma databricks.connection opção também se aplica a leituras e gravações de streaming quando sua fonte de dados implementa os métodos correspondentes (streamReader ou simpleStreamReader para leituras de streaming e writer ou streamWriter gravações).

Recursos adicionais

A comunidade Apache Spark mantém conectores de exemplo que você pode usar como referência ao construir suas próprias fontes de dados. Esses repositórios são mantidos pela comunidade e não suportados pelo Databricks:

• pyspark-data-sources: Uma coleção de exemplos de conectores de fonte de dados personalizados do PySpark.
• pyspark_huggingface: Um conector de fonte de dados personalizado para leitura de conjuntos de dados do Hugging Face.

Solução de problemas

Se a saída for o seguinte erro, sua computação não oferecerá suporte a fontes de dados personalizadas do PySpark. Você precisa usar o Databricks Runtime 15.2 ou superior.

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