PySpark カスタム データ ソース

PySpark カスタム データ ソースは 、Python (PySpark) DataSource API を使用して作成されます。これにより、Python を使用してカスタム データ ソースから読み取り、Apache Spark のカスタム データ シンクに書き込むことができます。 PySpark カスタム データ ソースを使用して、データ システムへのカスタム接続を定義し、再利用可能なデータ ソースを構築するための追加機能を実装できます。

SparkはDelta、Iceberg、Parquet、JSON、CSV、JDBCなどの標準フォーマットを組み込みサポートしていますが、REST API、Google Sheets、Hugging Faceデータセット、独自社内サービスなど多くのシステムには対応していません。 Python DataSource APIはこのギャップを埋めます。これらのシステムへのコネクターは純粋なPythonで構築し、JVMベースのコネクタ開発は不要で、Spark SQLを含む組み込みのSparkデータソースと同様に使用できます。

DataSource クラス

PySpark DataSource は、データ リーダーとライターを作成するメソッドを提供する基底クラスです。

データ ソース サブクラスを実装する

ユース ケースに応じて、データ ソースを読み取り可能か書き込み可能、あるいはその両方にするためには、何らかのサブクラスで実装する必要があります。

データ ソースを登録する

インターフェイスを実装した後、登録してから、次の例に示すように読み込むか、使用できます。

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

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

例 1: バッチ クエリ用の PySpark DataSource を作成する

PySpark DataSource のリーダー機能のデモを行うためにfaker Python パッケージを使用して、サンプル データを生成するデータ ソースを作成します。 faker に関する詳細については、「Faker ドキュメント」を参照してください。

次のコマンドを使用して、faker パッケージをインストールします。

%pip install faker

手順 1: バッチ クエリのリーダーを実装する

まず、リーダー ロジックを実装してサンプル データを生成します。 インストールされている faker ライブラリを使用して、スキーマ内の各フィールドを設定します。

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)

手順 2: DataSource の例を定義する

次に、新しい PySpark DataSource を、名前、スキーマ、およびリーダーを持つ DataSource のサブクラスとして定義します。 バッチ クエリでデータ ソースからの読み取りを行うには reader() メソッドを定義する必要があります。

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)

手順 3: サンプル データ ソースを登録して使用する

このデータ ソースを使用するために、その登録を行います。 既定では、FakeDataSourceには 3 つの行があり、スキーマには string、name、date、zipcodeのstate フィールドが含まれます。 次の例では、既定値を使用してサンプル データ ソースの登録、読み込み、出力を行います。

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

stringフィールドのみがサポートされていますが、faker パッケージ プロバイダーのフィールドに対応する任意のフィールドでスキーマを指定して、テストおよび開発用のランダム データを生成できます。 次の例では、name および 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   |
+---------------------+--------------+

カスタムの行数でデータ ソースを読み込むには、numRows オプションを指定します。 次の例では、5 個の行を指定します。

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

例2:バッチクエリでカスタムデータシンクに書き込み

PySpark DataSourceライターの機能を示すために、DataFrameの各パーティションをファイルに書き込み、ジョブがコミットした際に要約マーカーファイルを書き込むデータソースを作成します。

ステップ1:バッチクエリ用のライターを実装する

まず、ライターの論理を実装します。 各エグゼキューターはパーティションごとに1回 write() 呼び出しを行います。 すべての書き込みタスクが成功すると、ドライバーは commit()を呼び出します。 もし作業が失敗した場合、ドライバーが代わりに呼び abort() します。

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

ステップ2:書き込み可能なデータソースを定義する

次に、writer()を実装するDataSourceサブクラスを定義します。 overwrite 引数は、書き込みモードが True の場合は overwrite であり、False の場合は 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)

ステップ3:レジスタおよびデータシンクへの書き込み

このデータ ソースを使用するために、その登録を行います。 次に、ショートネームを format() に渡し、出力ディレクトリを path オプションに渡してDataFrameを書きます。 この例はUnityカタログのボリューム内のパスに書き込みます。 <catalog>、<schema>、<volume>を既存のボリュームに置き換えます。

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

例3:バリアントを使ってPySpark GitHub DataSourceを作成する

PySpark DataSource でのバリアントの使用を示すために、この例では GitHub からプル要求を読み取るデータ ソースを作成します。

バリアントの詳細については、「 クエリバリアントデータ」を参照してください。

手順 1: プル要求を取得するリーダーを実装する

まず、リーダー ロジックを実装して、指定した 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")
            )

手順 2: GitHub DataSource を定義する

次に、新しい PySpark GitHub DataSource を、名前、スキーマ、メソッドDataSourceを持つreader()のサブクラスとして定義します。 スキーマには、 id、 title、 user、 created_at、 updated_atの各フィールドが含まれます。 user フィールドはバリアントとして定義されます。

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)

手順 3: データ ソースを登録して使用する

このデータ ソースを使用するために、その登録を行います。 次の例では、データ ソースを登録して読み込み、GitHub リポジトリ PR データの 3 行を出力します。

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

例4:ストリーミングの読み書き用のPySpark DataSourceを作成する

PySpark DataSource ストリーム リーダーおよびライター機能のデモを行うために、faker Python パッケージを使用してすべてのマイクロバッチに 2 行を生成するサンプル データ ソースを作成します。 faker に関する詳細については、「Faker ドキュメント」を参照してください。

次のコマンドを使用して、faker パッケージをインストールします。

%pip install faker

手順 1: ストリーム リーダーを実装する

まず、すべてのマイクロバッチに 2 つの行を生成するストリーミング データ リーダーの例を実装します。 DataSourceStreamReaderを実装することも、データ ソースのスループットが低く、パーティション分割を必要としない場合は、代わりにSimpleDataSourceStreamReaderを実装できます。 simpleStreamReader() または streamReader() を実装する必要があり、simpleStreamReader() は streamReader() が実装されていないときにのみ呼び出されます。

DataSourceStreamReader の実装

streamReader インスタンスには、DataSourceStreamReader インターフェイスで実装されるすべてのマイクロバッチで 2 ずつ増加する整数オフセットがあります。

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

SimpleDataSourceStreamReader の実装

SimpleStreamReader インスタンスは、すべてのバッチに 2 行を生成する FakeStreamReader インスタンスと同じですが、パーティション分割なしで SimpleDataSourceStreamReader インターフェイスで実装されます。

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

手順 2: ストリーム ライターを実装する

次に、ストリーミング ライターを実装します。 このストリーミング データ ライターは、各マイクロバッチのメタデータ情報をローカル パスに書き込みます。

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

手順 3: DataSource の例を定義する

次に、新しい PySpark DataSource を、名前、スキーマ、およびメソッドのDataSourceとstreamReader()を持つ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)

手順 4: サンプル データ ソースを登録して使用する

このデータ ソースを使用するために、その登録を行います。 登録後、短い名前またはフル ネームを format()に渡すことで、ストリーミング クエリでソースまたはシンクとして使用できます。 次の例では、データ ソースを登録し、サンプル データ ソースから読み取ってコンソールに出力するクエリを開始します。

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

または、次のコードでは、ストリームの例をシンクとして使用し、出力パスを指定します。

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

例5:Google BigQueryストリーミングコネクターの作成

次の例では、PySpark DataSource を使用して Google BigQuery (BQ) 用のカスタム ストリーミング コネクタを構築する方法を示します。 Databricks には、BigQuery バッチ インジェスト用 の Spark コネクタ が用意されています。 また、Lakehouse Federation は、外部カタログの作成を通じて任意の BigQuery データ セットにリモート接続してデータをプルすることもできますが、増分または継続的ストリーミング ワークフローは完全にはサポートされていません。 このコネクタにより、段階的な増分データ移行と、永続的なチェックポイント処理を使用したストリーミング ソースによって提供される BigQuery テーブルからのほぼリアルタイムの移行が可能になります。

このカスタム コネクタには、次の機能があります。

• 構造化ストリーミングおよび Lakeflow パイプラインと互換性があります。
• 増分レコード追跡と継続的ストリーミング インジェストをサポートし、構造化ストリーミング セマンティクスに従います。
• より高速で安価なデータ転送のために、RPC ベースのプロトコルで BigQuery Storage API を使用します。
• 移行されたテーブルを Unity カタログに直接書き込みます。
• 日付またはタイムスタンプベースの増分フィールドを使用して、チェックポイントを自動的に管理します。
• Trigger.AvailableNow()でのバッチ インジェストをサポートします。
• 中間クラウド ストレージは必要ありません。
• Arrow または Avro 形式を使用して BigQuery データをシリアル化します。
• 自動並列処理を処理し、データ ボリュームに基づいて Spark worker 間で作業を分散します。
• SCD タイプ 1 またはタイプ 2 パターンを使用した Silver および Gold レイヤーの移行をサポートする BigQuery からの Raw および Bronze レイヤーの移行に適しています。

前提条件

カスタム コネクタを実装する前に、必要なパッケージをインストールします。

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

手順 1: ストリーム リーダーを実装する

まず、ストリーミング データ リーダーを実装します。 DataSourceStreamReader サブクラスは、次のメソッドを実装する必要があります。

• 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

各メソッドの詳細については、「 メソッド」を参照してください。

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

手順 2: データソースを定義する

次に、カスタム データ ソースを定義します。 DataSource サブクラスは、次のメソッドを実装する必要があります。

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

各メソッドの詳細については、「 メソッド」を参照してください。

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)

手順 3: ストリーミング クエリを構成して開始する

最後に、コネクタを登録し、ストリーミング クエリを構成して開始します。

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

実行順序

カスタム ストリームのファンクトン実行順序を以下に示します。

Spark ストリーム DataFrame を読み込む場合:

name(cls)
schema()

新しいクエリの開始または既存のクエリ (新規または既存のチェックポイント) の再起動時のマイクロバッチ (n) の場合:

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

既存のチェックポイントで実行中のクエリの次の (n+1) マイクロバッチの場合:

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

例6:外部APIで認証する

この例では、Unity Catalog の HTTP 接続を使用して外部 HTTP API に対して PySpark データ ソースを認証し、データ ソースのコードにハードコードされたトークンや認証情報が含まれないようにする方法を示します。

手順 1: HTTP 接続を作成する

データ ソースを実装する前に、Unity カタログに my_weather_api という名前の HTTP 接続を作成し、ユーザーまたはグループにアクセス許可 MANAGE 付与します。 接続に対する MANAGE アクセス許可を持つユーザーのみが資格情報の挿入をトリガーできます。

API トークンを Databricks シークレットとして格納し、リテラル トークンを入力するのではなく、 secret 関数で参照します。そのため、資格情報は接続定義に表示されません。

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

手順 2: バッチ クエリのリーダーを実装する

次に、REST API からデータをフェッチするリーダー ロジックを実装します。 リーダーは、挿入された host、 base_path、および bearer_token 値をオプションから読み取るので、コードに資格情報は表示されません。

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

手順 3: DataSource の例を定義する

次に、新しい PySpark DataSource を、名前、スキーマ、およびリーダーを持つ DataSource のサブクラスとして定義します。

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)

手順 4: データ ソースを登録して使用する

このデータ ソースを使用するために、その登録を行います。 次に、 databricks.connection オプションを使用して Unity カタログ HTTP 接続を参照します。 Spark ドライバーは、有効期間の短い OAuth2 資格情報を Unity カタログから自動的に取得し、データ ソース オプション マップに (たとえば、 bearer_token、 host、 base_path) 挿入します。 Unity カタログによって挿入された資格情報キーをオーバーライドすることはできません。また、 host や portなど、グローバルにブロックされるオプションはブロックされたままであり、ユーザーが設定することはできません。

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

この例では、バッチ読み取りのみを実装します。 同じ databricks.connection オプションは、データ ソースが対応するメソッド (ストリーミング読み取りの場合はstreamReader または simpleStreamReader 、書き込みの場合は writer または streamWriter ) を実装する場合のストリーミング読み取りと書き込みにも適用されます。

その他のリソース

Apache Sparkコミュニティは、自分でデータソースを構築する際の参照として使えるサンプルコネクターを管理しています。 これらのリポジトリはコミュニティによって管理されており、Databricksによってサポートされていません:

• pyspark-data-sources:PySparkのカスタムデータソースコネクタの例集です。
• pyspark_huggingface:Hugging Faceデータセットを読むためのカスタムデータソースコネクターです。

トラブルシューティング

出力が次のようなエラーとなる場合、お使いのコンピューティングでは PySpark カスタム データ ソースがサポートされていません。 Databricks Runtime 15.2 以上を使用する必要があります。

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