0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

Databricks Spark宣言型パイプラインのUnit Testing (Beta)を試してみる

0
Posted at

はじめに

Databricksのリリースノートを眺めていたら、以下の機能リリースが記載されていました。

Unit testing for pipelines (Beta)
You can now write unit tests for pipelines using the TestPipeline class.

Spark宣言型パ\イプライン(Spark Declarative Pipelines: 以下SDP)のUnit Testing機能です。
SDPの単体テストは以前からUIに出ていたように思うのですが、ベータとして試用できる状態になったということかな?
個人的にはパイプラインの単体テストが実行しやすくなるのは非常にありがたくです。そこで、公式ドキュメントに従って実際に試してみました。

この記事で扱う機能は2026年8月現在 Beta版であり、今後仕様が変更される可能性があります。

Spark宣言型パ\イプラインのUnit Testingとは

SDPのUnit Testingは、パイプラインのデータセット定義をpytestで検証できる機能です。

主な特徴は以下の通りです:

  • TestPipelineクラス: パイプラインをテスト環境で実行するためのクラス
  • test_sparkフィクスチャ: テスト用のSparkSessionを提供するpytestフィクスチャ
  • モックデータの作成: ソーステーブルをテスト用に作成・挿入してテスト実行
  • 出力検証: パイプライン実行後の出力テーブルを検証

従来、SDP自体のパイプラインテストはパイプラインを個別実行して結果を確認するしかありませんでしたが、この機能により個別のデータセットを使った再現可能な単体テストをできるようになりました。

検証環境

今回はDatabricks Free Editionを検証に利用しています。
有償版でも同様かと思います。

今回作るもの

DABを使ったシンプルなSDPパイプラインとそのUnit Testを作成します。

具体的には:

  1. パイプライン: NYC Taxiサンプルデータを読み込むマテリアライズドビュー
  2. Unit Test: モックデータを使ったパイプラインの動作検証

構成は以下のようになります:

test_bundle/
├── databricks.yml          # DABs設定
├── resources/
│   └── test_bundle_etl.pipeline.yml  # パイプライン定義
└── src/test_bundle_etl/
    ├── transformations/
    │   └── sample_trips.py    # パイプライン実装
    └── tests/
        └── test_trip.py       # Unit Test

Step0: プロジェクト準備

Declarative Automation Bundles(DABs)を使ってプロジェクトを作成します。
お手軽にUIから「Lakeflow Pipelines」のテンプレートを用いて作成します。

image.png

Step1: パイプラインの作成

まず、テスト対象となるシンプルなパイプラインを作成します。
Step0のテンプレートを使うと最初から作成されますが、以下のようなNYC Taxiのサンプルデータを読み込むだけのマテリアライズドビューです。

sample_trips.py
from pyspark import pipelines as dp
from pyspark.sql.functions import col

@dp.table
def sample_trips():
    return spark.read.table("samples.nyctaxi.trips")

パイプライン設定は resources/test_bundle_etl.pipeline.yml に定義があります。

test_bundle_etl.pipeline.yml
resources:
  pipelines:
    test_bundle_etl:
      name: test_bundle_etl
      target: ${var.catalog}.${var.schema}
      catalog: ${var.catalog}
      libraries:
        - glob:
            include: "/Workspace${workspace.file_path}/src/test_bundle_etl/transformations/**"
      continuous: false
      channel: PREVIEW
      serverless: true

continuous: falsechannel: PREVIEW を指定することで、Unit TestingのBeta機能が有効になります。
その設定を追加してデプロイしておきます。

Step2: テストコードの作成

次に、Unit Testを作成します。
パイプラインのメニューからテストを追加することでtestsフォルダとテスト用pythonファイルが作成されます。

image.png

作成したファイルにテストコードを記述していきます。
まず、必要なインポートとTestPipelineの初期化です:

test_trip.py
import pytest
from pyspark.pipelines.testing import TestPipeline, test_spark

test_pipeline = TestPipeline.active()

TestPipeline.active() は、現在のパイプラインコンテキストを取得します。

次に、ヘルパー関数を定義します。
現時点でパイプラインのデフォルトカタログ・スキーマの設定を取得する方法がわからなかったため、パイプライン実行結果のevent_logテーブルのパスからカタログ・スキーマ名を抽出しています。

test_trip.py
_CATALOG = "workspace"
_SCHEMA = "xxxxx"

def parse_catalog_schema_from_event_log_table(event_log_table_name: str | None):
    """イベントログテーブル名からcatalogとschemaを抽出"""
    if not event_log_table_name:
        return _CATALOG, _SCHEMA

    parts = event_log_table_name.split(".")
    if len(parts) != 3:
        raise ValueError("テーブル名は 'catalog.schema.table' フォーマットである必要があります")
    catalog, schema, _ = parts
    return catalog, schema

そして、メインのテストケースです。まずは正常系のテスト:

test_sample_trips_reads_trips_data(長いので折り畳み)
test_trip.py
def test_sample_trips_reads_trips_data(test_spark):
    """sample_tripsがtripsテーブルから正しくデータを読み込むことを検証"""

    # モックのソーステーブルを作成
    test_spark.sql("""
        CREATE OR REPLACE TABLE samples.nyctaxi.trips (
            tpep_pickup_datetime TIMESTAMP,
            tpep_dropoff_datetime TIMESTAMP,
            trip_distance DOUBLE,
            fare_amount DOUBLE,
            pickup_zip INT,
            dropoff_zip INT
        )
    """)
    
    # サンプルデータを挿入
    test_spark.sql("""
        INSERT INTO samples.nyctaxi.trips VALUES
        ('2025-01-01T08:00:00', '2025-01-01T08:15:00', 2.5, 12.0, 10001, 10002),
        ('2025-01-01T09:00:00', '2025-01-01T09:30:00', 5.0, 20.0, 10003, 10004),
        ('2025-01-01T10:00:00', '2025-01-01T10:10:00', 1.2, 8.5, 10005, 10006)
    """)

    # パイプラインを実行
    status = test_pipeline.run(test_spark, set(["sample_trips"]))
    
    # パイプライン実行が成功したことを確認
    assert status.is_success, f"Pipeline run failed: {status.error_message}"
    
    # 出力テーブルを検証
    catalog, schema = parse_catalog_schema_from_event_log_table(status.event_log_table_name)    
    output_table = f"{catalog}.{schema}.sample_trips"
    result = test_spark.table(output_table)
    
    # 行数を確認
    assert result.count() == 3, "Expected 3 trips in the output"
    
    # データ内容を確認
    rows = result.collect()
    
    # 1行目の検証
    assert rows[0].trip_distance == 2.5
    assert rows[0].fare_amount == 12.0
    assert rows[0].pickup_zip == 10001
    
    # 2行目の検証
    assert rows[1].trip_distance == 5.0
    assert rows[1].fare_amount == 20.0
    
    # 3行目の検証
    assert rows[2].trip_distance == 1.2
    assert rows[2].fare_amount == 8.5

テストの流れは以下の通りです:

  1. モックテーブルの作成: CREATE OR REPLACE TABLE でソーステーブルを作成
  2. テストデータ投入: INSERT INTO でサンプルデータを挿入
  3. パイプライン実行: test_pipeline.run() で指定したデータセットを実行
  4. 結果検証: 出力テーブルの内容をアサート

次に、エッジケースのテストです。空のソーステーブルでもパイプラインが正常に動作することを確認します:

test_sample_trips_empty_source(長いので折り畳み)
test_trip.py
def test_sample_trips_empty_source(test_spark):
    """sample_tripsが空のソーステーブルを正しく処理することを検証"""
    # 空のモックテーブルを作成
    test_spark.sql("""
        CREATE OR REPLACE TABLE samples.nyctaxi.trips (
            tpep_pickup_datetime TIMESTAMP,
            tpep_dropoff_datetime TIMESTAMP,
            trip_distance DOUBLE,
            fare_amount DOUBLE,
            pickup_zip INT,
            dropoff_zip INT
        )
    """)
    
    # パイプライン実行
    status = test_pipeline.run(test_spark, set(["sample_trips"]))
    
    # 空の入力でも成功することを確認
    assert status.is_success, f"Pipeline run failed: {status.error_message}"
    
    catalog, schema = parse_catalog_schema_from_event_log_table(status.event_log_table_name)    
    output_table = f"{catalog}.{schema}.sample_trips"
    result = test_spark.table(output_table)
    
    # 出力も空であることを確認
    assert result.count() == 0, "Expected empty output for empty input"

Step3: テストの実行

テストを実行します。
Databricks Pipeline Editor上に、テスト用のアイコンメニューがありますので、それを開いて実行します。

image.png

テスト結果はエディタ下部、出力やターミナル等と同じタブのところに表示されます。
2つのテストケースがともにPASSEDとなり、パイプラインが正しく動作していることが確認できました。実行時間は約66秒でした。

image.png

詳細な実行ログや、エラーがある場合はエラーログを表示できます。

ハマりポイント・注意点

実際に試してみて、いくつかハマったポイントがありましたので共有します。

1. パイプラインの設定情報取得

テスト内で、パイプラインのデフォルトカタログ名・スキーマ名を取得したかったのですが、方法が結局わかりませんでした。
(今回は、無理やりテスト内でのパイプライン処理実行後、status.event_log_table_name から出力テーブルのcatalog/schemaを取得しています。あまりよいやり方ではないです)

テストにおいては、テーブルのパスをcatalog.schema.tableのフルで指定するべきとは思うのですが、パイプラインともう少し連動する機能が欲しいと思いました。

2. channelの指定

Beta機能を使うには、パイプライン設定で channel: PREVIEW を指定する必要があります。
現時点では、これを忘れるとUIからテストを実行できません。
ドキュメントにも明記されているのでよく読みましょう。

まとめ

Spark宣言型パイプラインのUnit Testing機能(Beta)を試してみました。
テストのベース作成はGenie Codeに任せていけますし、ドキュメントにはより有用なテスト例も記載されているのでSDPを利用しているプロジェクトをされている方は検証していってよい機能かなと思います。

ただし、まだBeta版ということもあってか、パイプラインとの連動機能がまだ少ないとか、割と実行時間が長く開発体験がイマイチとか、まだまだこれからな印象です。
私はSDPをよく使うので、この機能の動向は引き続き追っていきたいと思います。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?