Claude Media
Claude CodeでPySparkのジョブを書く — local[2]のpytestで検証する手順

Claude CodeでPySparkのジョブを書く — local[2]のpytestで検証する手順

PySparkの変換ロジックをClaude Codeに書かせるとき、local[2]のSparkSession fixtureとassertDataFrameEqualで合否を決める手順と、CLAUDE.mdと権限の設定をまとめます。

PySparkのジョブをClaude Codeに書かせるなら、完了の合図はクラスターでの実行ではなく、手元のpytestの合否に置くのが扱いやすい構成です。local[2]でSparkSessionを立て、変換結果をassertDataFrameEqualで期待値と突き合わせます。クラスターもクラウドの認証も要らず、Claudeは編集のたびに同じコマンドで結果を読めます。この記事では、テストしやすいジョブの形、fixture、比較関数の引数、CLAUDE.mdと権限の設定までを順に示します。

前提: 用意するもの

手元に必要なのは、PythonとJava、2つのPythonパッケージです。

  • pyspark: Pythonのパッケージ。pip install pysparkで入ります
  • pytest: テストランナー。pip install pytestで入ります
  • Java: PySparkはJava 17以降を求め、環境変数JAVA_HOMEが正しく設定されている必要があります
  • Spark 3.5以降: Sparkのテストガイドの例はこのバージョンが対象です

PySparkの導入ページに書かれているのは「Java 17以降」までで、そのページの内容はlatest版です。使うPySparkのバージョンによって対応するJavaが変わることがあるため、古いバージョンを固定している場合は、そのバージョンの導入ページで確かめてください。

pip install pyspark pytest
java -version

java -versionが通らない環境では、SparkSessionの起動で失敗します。Claudeにテストを回させる前に、人が一度この2行を通しておくと、原因の切り分けが楽です。

検証コマンドをpytest1本に決める

Claude Codeは「テストが通った」という具体的な出力を得ると、次の修正の向きを自分で決められます。PySparkでその役を担うのが、ローカルモードのSparkSessionに対するpytestです。

Sparkのテストガイドは、PySparkの組み込みテスト関数をunittestやpytestなど任意のテストフレームワークで使えると説明しています。関数は単体で動くため、CI(継続的インテグレーション)にもそのまま載せられます。使うSparkのバージョンが古い場合は、まずpyspark.testingが入っているかを確かめてください。

検証コマンドは次の1本に絞ります。

python -m pytest tests -q

全体を回しても数秒から数十秒で終わる粒度を保つと、Claudeに毎回実行させても待ち時間が負担になりません。遅くなる原因の多くはSparkSessionの起動回数とシャッフルの分割数です。後半で順に扱います。

テストしやすいジョブの形にする

ジョブ全体を1つのスクリプトに書くと、読み込みと書き出しにテストが引きずられます。変換だけを「DataFrameを受け取ってDataFrameを返す関数」に切り出します。SparkSessionは関数の外で作って渡す形にします。

repo/
├── pkg/
│   └── etl.py          # 変換関数だけを置く
├── tests/
│   ├── conftest.py     # SparkSession fixture
│   └── test_etl.py
└── CLAUDE.md

ガイドも、空白を正規化する関数remove_extra_spaces(df, column_name)を定義し、その関数をテストする構成です。この形に合わせた変換関数の例を示します。

# pkg/etl.py
from pyspark.sql import DataFrame
from pyspark.sql.functions import col, regexp_replace, sum as sum_
 
 
def remove_extra_spaces(df: DataFrame, column_name: str) -> DataFrame:
    return df.withColumn(
        column_name, regexp_replace(col(column_name), r"\s+", " ")
    )
 
 
def total_by_id(df: DataFrame) -> DataFrame:
    return df.groupBy("id").agg(sum_("amount").alias("total"))

2つ目のtotal_by_idはgroupByを使うため、実行時にシャッフルが発生します。fixtureの設定が効く場面を見せるために入れました。

local[2]のSparkSession fixtureを作る

conftest.pyに置いたfixtureは、tests配下のすべてのテストから名前だけで使えます。ここでSparkSessionのマスターをlocal[2]に固定します。local[K]はK個のワーカースレッドでローカル実行するモードです。localだけだと並列度は1になります。local[2]なら最小限の並列性を保ったまま、ローカルのコア数に依存しません。

# tests/conftest.py
import pytest
from pyspark.sql import SparkSession
 
 
@pytest.fixture(scope="session")
def spark():
    session = (
        SparkSession.builder.master("local[2]")
        .appName("pytest-pyspark")
        .config("spark.sql.shuffle.partitions", "2")
        .config("spark.ui.enabled", "false")
        .getOrCreate()
    )
    yield session
    session.stop()

設定は4つです。

  • master("local[2]"): クラスターを使わず2スレッドで実行する
  • scope="session": テスト全体で1つのSparkSessionを共有する
  • spark.sql.shuffle.partitionsを2にする: 既定値は200で、小さなテストデータには過剰な分割になる
  • spark.ui.enabledをfalseにする: Web UIの起動を省く。既定はtrue

Sparkのガイドにあるpytestの例は、SparkSession.builder.appName(...).getOrCreate()を呼んでyieldするだけの関数スコープのfixtureです。マスターの指定もstop()もありません。そのままだと実行環境の既定設定に依存し、テストごとにセッションを作り直す形にもなります。上の形に直しておくと、Claudeが新しいテストを足すときも同じfixtureを再利用します。

assertDataFrameEqualで期待値と照合する

DataFrameの比較はpyspark.testingのassertDataFrameEqualを使います。2つのDataFrameを受け取り、違いがあれば例外を投げます。シグネチャは次のとおりです。

assertDataFrameEqual(
    actual, expected,
    checkRowOrder=False, rtol=1e-05, atol=1e-08,
    ignoreNullable=True, ignoreColumnOrder=False,
    ignoreColumnName=False, ignoreColumnType=False,
    maxErrors=None, showOnlyDiff=False, includeDiffRows=False,
)
引数既定値効き方
checkRowOrder既定値False効き方行の順序を比較しない。順序を保証したい出力だけTrueにする
rtol / atol既定値1e-05 / 1e-08効き方浮動小数点の相対・絶対許容誤差
ignoreNullable既定値True効き方スキーマのnullable属性の違いを無視する
ignoreColumnOrder既定値False効き方列の並び順の違いも差として扱う
ignoreColumnType既定値False効き方列の型の違いも差として扱う
maxErrors既定値None効き方行の比較失敗がこの件数に達した時点で比較を打ち切る。Noneは全行を比較する
showOnlyDiff既定値False効き方Trueだと、エラーメッセージに差のある行だけを載せる
includeDiffRows既定値False効き方Trueだと、差のある行をPySparkAssertionErrorに含めて後から調べられる

下の3行は、導入が4.0.0からです。Spark 3.5系のAPIページにはこの3つの引数が載っていないため、3.5を使う場合は上の5行だけが使えます。

showOnlyDiffとmaxErrorsは、Claudeが失敗出力を読むときに効きます。既定のままだと、1行でも違えばエラーメッセージに全行が並び、数百行のテストデータで出力が長くなります。showOnlyDiff=Trueで差のある行だけに絞ると、Claudeが読む量が減ります。差が大量に出る壊れ方では、maxErrors=5のように上限を置いて早めに打ち切ることもできます。

Sparkのガイドは、rtol=1e-1を渡して「0.1と0.109」を近似一致として通す例があります。許容誤差を広げるほど、Claudeが誤った計算を書いてもテストが通りやすくなります。金額や件数のように厳密であるべき列では、既定のrtolとatolのまま使います。

テストを書く

期待値は、仕様から人間が決めるか、仕様の文章をClaudeに渡して書かせたうえで人が確認します。実装を読んだClaudeが期待値を書くと、実装の誤りを写し取ってしまうことがあるためです。

# tests/test_etl.py
from pyspark.sql.types import (
    LongType, StringType, StructField, StructType,
)
from pyspark.testing import assertDataFrameEqual
 
from pkg.etl import remove_extra_spaces, total_by_id
 
 
def test_remove_extra_spaces(spark):
    original = spark.createDataFrame(
        [("John    D.", 30), ("Alice   G.", 25)], ["name", "age"]
    )
    expected = spark.createDataFrame(
        [("John D.", 30), ("Alice G.", 25)], ["name", "age"]
    )
    assertDataFrameEqual(remove_extra_spaces(original, "name"), expected)
 
 
def test_total_by_id(spark):
    schema = StructType([
        StructField("id", StringType()),
        StructField("amount", LongType()),
    ])
    original = spark.createDataFrame(
        [("a", 100), ("a", 50), ("b", 30)], schema
    )
    expected = spark.createDataFrame(
        [("a", 150), ("b", 30)], ["id", "total"]
    )
    assertDataFrameEqual(total_by_id(original), expected)

スキーマを明示した入力は、型の取り違えを早く見つけられます。ignoreColumnTypeの既定がFalseなので、期待値の型が違えばテストが落ちます。落ちたときのメッセージはClaudeがそのまま読み、関数側か期待値側のどちらを直すかを判断します。

実行を依頼するときの指示は短くて足ります。

pkg/etl.py に、重複行を除く dedupe_orders(df) を追加して。
先に tests/test_etl.py へ失敗するテストを書き、
python -m pytest tests -q が通るまで直して。
期待値の変更は私に確認してから。

「先に失敗するテストを書く」「期待値の変更は確認する」の2点が、このやり方の要です。

CLAUDE.mdに書く規約

プロジェクトのCLAUDE.mdに、検証コマンドと守る約束を置きます。pytest全般の書き方はCLAUDE.mdへの規約の入れ方を扱ったClaude CodeにpytestのCLAUDE.md規約を教えるにまとまっています。PySpark固有の部分だけを足します。

## PySpark
- 変換ロジックは pkg/ に「DataFrame を受けて返す関数」で書く。読み書きは分ける
- 検証は `python -m pytest tests -q` だけ。クラスターやクラウドへは接続しない
- SparkSession は tests/conftest.py の `spark` fixture を使う。テスト内で作らない
- DataFrame の比較は assertDataFrameEqual。rtol を広げない
- 期待値の変更は人の確認を得てから行う

テスト実行だけを許可して本番系を止める

CLAUDE.mdは指示であり、動作を強制しません。権限ルールはClaude Code側が強制します。Claude Codeの権限は、deny、ask、allowの順に評価され、先に一致したものが結果を決めます。広いdenyルールは、より狭いallowルールがあっても優先されます。

テストコマンドだけを許可し、spark-submitを拒否する設定の例です。

{
  "permissions": {
    "allow": [
      "Bash(python -m pytest *)"
    ],
    "deny": [
      "Bash(spark-submit *)",
      "Bash(databricks *)"
    ]
  }
}

ローカルの検証ではクラスターに送る必要がないため、spark-submitを最初から止めておいても支障がありません。設定ファイルに書いたルールは、プロンプトでの依頼やCLAUDE.mdの指示と違い、Claude Code側が実行を止めます。

テストを編集のたびに自動で回したい場合は、フックで実行を仕込めます。「通るまで止めない」型はClaude Codeのagent hookで「テストが通るまで止めない」を検証させるで扱っています。

つまずきやすい箇所

公式の例をそのまま写すと動かない箇所がある

ガイドのunittestの例では、setUpClassでcls.sparkを作りながら、テスト本体はspark.createDataFrameとグローバル変数のsparkを呼んでいます。ノートブックで前のセルが残っている前提の書き方です。pytestに移すときは、必ずfixtureの引数で受け取ります。Claudeに公式の例を渡すときも、この違いを指示に入れておきます。

テストが遅い

遅さの原因は、ほぼ次の3つです。

  1. テスト関数ごとにSparkSessionを作っている。scope="session"にする
  2. spark.sql.shuffle.partitionsが既定の200のまま。2程度に下げる
  3. 入力データが大きすぎる。数行から数十行の手書きデータで足りる

行の順序で落ちる

checkRowOrderの既定はFalseなので、順序の差では落ちません。順序が必要な出力だけ、orderByで並べたうえでTrueにします。

型が合わない

createDataFrameにタプルのリストを渡すと、Pythonの整数がLongTypeに推論されます。期待値のスキーマを明示すると、IntegerTypeかLongTypeかの食い違いに悩まずに済みます。

どちらの型で落ちているのか分からないときは、中身の比較より先にスキーマだけを照合します。pyspark.testingのassertSchemaEqualは、2つのスキーマを比べて違いがあれば例外を投げます。Sparkのテストガイドにも、StructType同士を渡す例があります。

from pyspark.testing import assertSchemaEqual
 
def test_total_by_id_schema(spark):
    original = spark.createDataFrame(
        [("a", 100)],
        StructType([
            StructField("id", StringType()),
            StructField("amount", LongType()),
        ]),
    )
    expected_schema = StructType([
        StructField("id", StringType()),
        StructField("total", LongType()),
    ])
    assertSchemaEqual(total_by_id(original).schema, expected_schema)

スキーマが合ってからassertDataFrameEqualで値を見る順にすると、Claudeに渡る失敗の原因が型か値かで分かれます。

Databricksなどクラスターへの展開

ここまでの検証は、ローカルで変換ロジックだけを確かめる範囲です。クラスターでの実行、ジョブの定義、デプロイは別の問題になります。Databricksへの展開はClaude CodeでDatabricks Asset Bundlesを書くが扱っています。dbtのようにウェアハウス側のSQLを検証する場合は、Claude Codeでdbt build --selectを回すの手順が参考になります。

まとめ

PySparkをClaude Codeに任せる土台は、local[2]のfixtureとassertDataFrameEqualによる「ローカルで完結する合否」です。変換を関数に切り出し、shuffle.partitionsを絞り、期待値の変更は人が確認する。この3つを守れば、Claudeの修正が正しいかどうかを毎回同じ基準で判定できます。

この記事を共有:XはてブLinkedIn