PySpark でデータを加工し、遅いコードを見抜く

5.1PySpark DataFrame API 実践 — select / filter / withColumn / when・otherwise / groupBy・agg のコード読解

Databricks の試験では「次のコードを実行すると何が返りますか」という設問が繰り返し出ます。手元で動かせない以上、コードを 1 行ずつ追って結果を言い当てる読み方を先に決めておくしかありません。読む手順は毎回同じで構いません。

1. アクションを探し、どこでジョブが起動するかを確かめる。
2. 行数が変わる箇所を探す (filter、groupBy、join、dropDuplicates)。
3. 列が増減する箇所を探す (select、withColumn、agg)。
4. 条件式の境界 (> か >= か) と NULL の扱いを確認する。
5. 選択肢のうち、行数と列名が合わないものを先に落とす。

コードは書いた行では動かない

select や filter は変換 (transformation) と呼ばれ、実行計画に積まれるだけで計算は始まりません。show() や count() のようなアクション (action) を呼んだ瞬間に、はじめてクラスタ上でジョブが起動します。

変換は元の DataFrame を書き換えません。4 行の df に df.filter(df.id > 1) と書いても df は 4 行のままで、戻り値を変数で受けなければ結果は捨てられます。「フィルタが効かない」の最も多い原因がこれです。次のコードで例外を捕まえられないのも、try ブロックにアクションが 1 つも無いためです。

try:
    df = spark.read.json("/path")
    df.withColumn("new", df["missing"] + 1)   # 変換だけ。ここでは走らない
except Exception as e:
    print("Error:", e)                        # ここには到達しない

select は列を選び、filter は行を残す

select は SQL の SELECT 句にあたり列を絞ります。filter は WHERE 句にあたり行を絞ります。where は filter の別名で、動作はまったく同じです。

df = spark.read.format("delta").load("/delta/sales")
result = df.select("product_id", "quantity").filter(df.quantity > 10)
result.show()

返るのは 2 列で、quantity が 10 を超える行だけです。> なので 10 ちょうどの行は残りません。境界を > と >= のどちらで書いているかは、コード読解問題で最も狙われる箇所です。

条件は df.filter(df.id > 1) のような Column 式と、df.filter("id > 1") のような SQL 文字列の 2 通りで書けます。ここで df.filter("amount" > 1000) と書くと、Python の文字列と数値を比べるだけになり行は絞られません。存在しない列名を渡した場合は AnalysisException: cannot resolve 'colX' になり、タイプミスか、JOIN 後に t1.colX と修飾すべき所を素で書いたか、のどちらかが原因です。

条件を 2 つ以上つなぐときは、Python の and / or ではなく & / | / ~ を使い、各条件を必ず丸括弧で囲みます。& は比較演算子より優先順位が高いため、括弧を省くと意図しない結合になって失敗します。

df.filter((col("qty") > 10) & (col("status") == "active"))   # 正しい
df.filter(col("qty") > 10 & col("status") == "active")       # 括弧が無く失敗する

NULL も罠です。filter(col("status") != "A") と書くと、status が NULL の行も落ちます。NULL との比較結果は真でも偽でもなく unknown で、unknown の行は残らないためです。残したいなら col("status").isNull() を | で足します。型のずれも例外なしに結果だけ変えるので、col("amount").cast("double") > 100 のように揃えてから比べます。

記述順序も問われます。行と列を先に削る filter → select → groupBy の順が読みやすい書き方です。ただし Catalyst Optimizer が述語プッシュダウン (条件を読み込み側へ前倒しする最適化) と列の刈り込みを自動で行うため、select と filter の前後は結果も物理プランも変わりません。「必ず filter を先に書かなければ正しく動かない」と断定する選択肢は誤りです。

withColumn は同名なら上書きする

withColumn(name, expr) は列を 1 本追加しますが、指定した名前が既存の列と同じなら、新しい式の結果で上書きされます。列が 2 本になることも、price_1 のような別名が自動で付くこともありません。

df.withColumn("price", col("price") * 1.1)   # price 列が 1.1 倍で置き換わる

公式ドキュメントは、このメソッドが内部で射影を 1 つ追加するため「繰り返し呼ぶと巨大なプランが生成され、性能問題や StackOverflowException を起こしうる」と警告しています。列を複数足すときは select に列式を並べるか、辞書を渡す withColumns({"a": expr1, "b": expr2}) で 1 回にまとめます。

when と otherwise は CASE WHEN そのもの

条件で値を出し分ける列は when().otherwise() で作ります。SQL の CASE WHEN ... THEN ... ELSE ... END と同じ意味で、上から順に評価して最初に真になった枝で確定します。どの条件にも当たらない行には otherwise の値が入り、otherwise が無ければ NULL になります。

df.withColumn("category",
  when(col("price") > 1000, "High")
  .when(col("price") > 300, "Mid")
  .otherwise("Low"))

閾値を広い順に並べた瞬間、後ろの狭い枝には誰も到達しません。設問では 1000 と 300 を入れ替えた選択肢が必ず混ざります。なおここで Python の if や三項演算子を使ってはいけません。when は各行で評価されますが、Python の if はドライバで 1 回だけ評価され、その結果が全行に同じ値として貼られます。

groupBy と agg は行を減らす

groupBy("store") が返すのは DataFrame ではなく GroupedData です。ここに agg() を続けると、グループごとの集計値を 1 回のシャッフル (同じキーの行を同じタスクへ集めるためのデータ再配置) で計算した DataFrame が得られます。

df.groupBy("store").agg(F.sum("amount").alias("total"), F.avg("amount"))  # 正しい
df.groupBy("store").agg({"amount": "sum"})                                # 辞書形式も可
df.groupBy("store").sum("amount").avg("amount")                           # AttributeError

3 行目が失敗するのは、GroupedData.sum() の戻り値が DataFrame で、DataFrame に avg というメソッドが無いためです。また F. を付けずに sum("amount") と書くと、Python 組込みの sum や col と名前が衝突して TypeError: 'Column' object is not callable が出ます。列名は sum(amount) の形になるので alias で付け直します。

NULL の扱いも設問になります。Databricks SQL では COUNT(*) を唯一の例外として、集計関数は NULL を無視します。一方 GROUP BY では複数の NULL が 1 つのグループにまとめられます。

値を配列にまとめる collect_list (別名 array_agg) も頻出です。公式ドキュメントは「NULL 値は除外される」「配列内の要素の順序は非決定的である」と明記しています。NULL を落としたいなら WHERE col IS NOT NULL で先に除き、代わりの値を残したいなら collect_list(coalesce(col, default)) と書きます。全行が NULL のグループは空配列になります。順序を確定させるには array_sort、重複を落とすには collect_set を使います。すでに配列型の列にもう一度 collect_list をかけると array<array<...>> の二重配列になる点にも注意します。

集約後にさらに絞るときは groupBy(...).agg(...).filter(...) と続けます。これは SQL の HAVING に相当し、集約前の filter とは意味が違います。そして集約対象を間違えても例外は出ません。売上列が amount なのに F.sum("qty") と書けば、返るのは在庫数の合計です。数値だけが静かにずれる誤りは設問でもこの形で出ます。

行と列を入れ替える explode と PIVOT

ウィンドウ関数は行を減らさずに集計値を付ける

groupBy が行を潰すのに対し、ウィンドウ関数は元の行数を保ったまま、各行に集計値やランクを列として足します。Window.partitionBy(...).orderBy(...) で範囲を決め、集計関数に .over(w) を付けます。

最頻出のイディオムが、キーごとに最新の 1 件だけを残す重複排除です。

w = Window.partitionBy("customer_id").orderBy(F.col("order_ts").desc())
latest = (df.withColumn("rn", F.row_number().over(w))
            .filter(F.col("rn") == 1))

ここを rank() にすると、同じ時刻の注文が 2 件ある顧客で rank = 1 の行が 2 行返り、1 行に絞れません。row_number は同点でも連番を振るので必ず 1 行になります。dense_rank は同点に同じ順位を与えつつ次の順位を飛ばさない関数で、こちらも重複排除には使えません。dropDuplicates(["customer_id"]) はどの行が残るか保証されないため、「最新の 1 件」という要件には row_number を使います。

累計を出すときはフレームを明示します。rowsBetween(Window.unboundedPreceding, Window.currentRow) が「先頭行から現在行まで」で、rowsBetween(0, Window.unboundedFollowing) と書くと現在行から末尾までになり、逆向きの残高集計に変わります。ROWS は物理的な行数、RANGE は ORDER BY 列の値が同じ行をひとまとめに扱う点が違います。

運用上の注意も 2 つ問われます。1 つ目は ORDER BY を省略する、あるいは順序が一意に決まらない列で並べると、row_number や lag の結果が実行のたびに変わることです。2 つ目は partitionBy の指定で、省略すると全行が 1 つのウィンドウにまとめられて 1 タスクに寄り、逆に偏ったキーを選ぶとそのキーだけが極端に重くなります。この偏りが第 3 節で扱うデータスキューです。

結果が 0 行、または想定と違うときの読み方

コードが動いたのに結果がおかしいときは、当てずっぽうに設定を変えず順番に確かめます。まず WHERE や filter をいったん外し、件数・スキーマ・サンプル行を見ます。これで「そもそもデータが無い」のか「条件で落としている」のかが 1 回で分かります。ウェアハウスのサイズを上げる、LIMIT 1 を付ける、といった対処は原因に届きません。中間の DataFrame は display() や printSchema() で逐次確認し、Databricks ノートブックなら変数エクスプローラで値と型を、組込みデバッガでブレークポイントとステップ実行を使えます。

5.2UDF は最後の手段 — 組込関数 > Pandas UDF > Python UDF > RDD.map という性能序列の理由

「Python で書けるから」と udf() で包んだ関数を数億行に適用した翌朝、10 分で終わっていたジョブが 3 時間かかっている。UDF の事故はだいたいこの形で表れます。厄介なのは例外が出ないことです。結果は正しく、ただ遅い。この節ではなぜ遅くなるのかを内部構造から押さえ、組込関数 > pandas UDF > Python UDF > RDD.map という序列を判断基準として使えるようにします。

Python UDF が遅い理由は 2 つある

1 つ目はプロセス間の往復です。実際の計算を担う Executor (ワーカーノード上で処理を実行するプロセス) は JVM で動きますが、Python UDF の本体は別の Python ワーカープロセスで動きます。素の Python UDF は行を 1 件ずつ pickle でシリアライズして Python 側へ送り、結果をデシリアライズして JVM へ戻します。この往復が行数の分だけ繰り返されます。

2 つ目は最適化から外れることです。Catalyst Optimizer は UDF の中身を解析できず、ブラックボックスとして扱います。述語プッシュダウンも定数畳み込みも、式を JVM のネイティブコードへまとめて生成する Whole-Stage Codegen も効きません。さらに Databricks の高速実行エンジンである Photon の公式ドキュメントは「Photon は UDF、RDD API、Dataset API をサポートしない」と明記しており、対応していない処理に当たると「その処理の残りについてコンピュートは透過的に Spark ランタイムへ切り替わる」と説明しています。結果は正しく返りますが、Photon を有効にしたクラスタの料金を払いながら恩恵だけが消えます。

置き換えの優先順位

まず疑うのは「そもそも UDF が要るのか」です。文字列なら upper や regexp_replace、時刻なら from_unixtime、分岐なら when があります。F.udf(lambda x: x + 1, IntegerType()) のような式は、F.col("v") + 1 と書けば JVM 内で完結します。

配列型の列を要素ごとに加工したいという理由で UDF を書くケースも多いのですが、これは高階関数で書けます。transform(items, x -> x * 1.1) は各要素を変換した同じ長さの配列を返し、filter(tags, t -> length(t) >= 3) は条件を満たす要素だけを残した配列を返します。aggregate(arr, 0, (acc, x) -> acc + x) は配列を 1 つの値に畳み込み、exists は 1 つでも条件を満たすかを真偽値で返します。高階関数は Catalyst の組込み式として JVM 内で完結するため、UDF のようなシリアライズ往復も、Photon が非対応と明記している UDF によるフォールバックも起こしません。ただし配列を扱う高階関数そのものは現行リリースの Spark では Whole-Stage Codegen の対象外です。UDF より確実に速い、というだけで、組込のスカラー関数と同じところまで最適化されるわけではありません。

pandas UDF の書き方と設定値

Python でしか書けないロジックが残ったら pandas UDF に切り替えます。公式ドキュメントは、行単位の Python UDF に対して最大 100 倍の性能向上があるとしています。

from pyspark.sql.functions import pandas_udf
import pandas as pd

@pandas_udf("double")
def sq(s: pd.Series) -> pd.Series:
    return s * s

df = df.withColumn("squared", sq(df["value"]))

覚えておく設定値は spark.sql.execution.arrow.maxRecordsPerBatch で、パーティションは既定で 1 バッチ 10,000 レコードずつ Arrow のレコードバッチに分割されて転送されます。メモリ不足が起きるときはこの値を下げます。戻り値の型は @pandas_udf("double") のように文字列で書くか型ヒントで示し、必ず明示します。

型は用途別に分かれます。1 行 1 値を返す Series から Series への型がベクトル化スカラー、モデルの読み込みなど初期化を 1 回で済ませたいときが Iterator を使う型、グループごとに 1 値へ集約するのが Series からスカラーへの型です。最後の型には注意が必要で、公式ドキュメントは部分集約に対応せず、各グループのデータがすべてメモリに読み込まれると明記しています。グループが巨大だとここで落ちます。なお pandas UDF が速いのは Arrow のバッチ転送と pandas のベクトル演算のおかげで、Python 部分は依然として Python プロセスで動きます。「GIL を回避して JVM 上でネイティブ実行される」という説明は誤りです。

やってはいけない対処

• repartition(1) で 1 タスクに集める。並列度が完全に失われ、むしろ悪化します。
• UDF を Python のグローバル変数にする。シリアライズの回数もコストも変わりません。
• UDF の中に print() を入れる。行数の分だけ I/O が増えます。
• UDF の中から外部 REST API を同期呼び出しする。1 行あたり数百ミリ秒の待ちが行数分積み上がります。

一方、UDF が一部の行で例外を投げてタスクごと落ちる場合は、想定できる例外だけを try / except で捕まえ、失敗した行に None を返すのが定石です。後段で NULL を検知すれば品質チェックに乗せられます。except Exception や BaseException で全部を握りつぶすとサイレント障害になるため、捕捉する例外型は具体的に指定します。

5.3JOIN が遅いとき何が起きているか — NULL キー脱落・型不一致・Broadcast Join・pushdown・AQE・skew・spill

JOIN の相談は 2 種類に割れます。行が減ったか、終わらないかです。混ぜて調べると、正しくない結果を速く出すだけの改善になります。まず行数、次に時間の順で切り分けます。クラスタを大きくする、cache() を足す、FULL OUTER JOIN に書き換える、といった対処はどちらの原因にも届かないので、選択肢に出てきたら真っ先に外して構いません。

行が減る — NULL キーと型不一致

INNER JOIN は結合キーが両側で一致した行だけを返します。ここで効くのが NULL の扱いです。Databricks の公式ドキュメントは通常の比較演算子について「オペランドの一方または両方が NULL のとき、結果は unknown または NULL になる」と定めています。つまり NULL = NULL は真になりません。キーに NULL が混ざった行は、警告も出ないまま静かに脱落します。

対処は前段に WHERE key IS NOT NULL を置いて落ちる行を可視化し、必要なら隔離用のテーブルへ退避することです。coalesce(key, 0) でゼロ埋めするのは、無関係な行同士を誤って結合させる危険な回避策で、集計値が静かに膨らみます。NULL 同士を一致させたい特殊な要件があるときだけ、NULL セーフ等価演算子 <=> (PySpark では eqNullSafe) を使います。この演算子は片方だけが NULL なら False、両方が NULL なら True を返します。

次に疑うのが型です。STRING の "001" と INT の 1、前後に混ざった空白、大文字小文字の違いはいずれも一致失敗を招きます。両側のキーについて件数・NULL 件数・distinct 件数・データ型を数えるところから始めます。

終わらない — どの結合戦略が選ばれたか

結合戦略は主に 2 つです。両側をキーでシャッフルして突き合わせる Sort Merge Join と、小さい側を全 Executor に配って各所でハッシュ照合する Broadcast Hash Join です。前者は全件シャッフルが起き、後者はシャッフルが要りません。10 億行のテーブルと 1 万行のマスタなら後者が圧倒的に速くなります。

サイズ推定が外れて自動判定が働かないときは、SELECT /*+ BROADCAST(r) */ ... というヒントで強制できます。ヒントは BROADCAST・MERGE・SHUFFLE_HASH・SHUFFLE_REPLICATE_NL の 4 種類で、両側に別々のヒントが付いた場合はこの並び順の優先度で解決されます。配るのは常に小さい側です。10 億行の側を指定すればドライバとエグゼキュータのメモリが破綻します。

AQE が実行時に直してくれる範囲

Adaptive Query Execution (AQE) は、シャッフルで得た実測の統計値を使って実行の途中でプランを組み替える仕組みです。Databricks では spark.databricks.optimizer.adaptive.enabled の既定が true で、最初から有効です。公式ドキュメントが挙げる働きは 4 つあります。Sort Merge Join から Broadcast Hash Join への動的な切り替え、シャッフル後のパーティションの動的な結合、Sort Merge Join と Shuffle Hash Join における偏りの動的な処理、そして空のリレーションの検出と伝播です。

偏りの判定条件は具体的に決まっています。あるパーティションが spark.sql.adaptive.skewJoin.skewedPartitionFactor (既定 5) を中央値に掛けた大きさを超え、かつ spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes (既定 256MB) を超えたときに、そのパーティションが偏っていると判定されます。両方を満たす必要がある点が問われます。データ量が小さくて 256MB に届かない偏りは AQE では直らないため、キーの見直しやソルティングが要ります。結合後のパーティションの大きさは spark.sql.adaptive.advisoryPartitionSizeInBytes (既定 64MB) が目安です。手動の SKEW ヒントも残っていますが、公式ドキュメントは「AQE のスキュー結合処理は完全に自動で、一般にヒントより性能が良いため、ヒントではなく AQE に任せることを推奨する」としています。

Spark UI で症状を読む

ここまでの原因は Spark UI のステージ画面で見分けられます。Databricks の公式ガイドは「スピルは Spark のメモリが足りなくなったときに起きる。データをメモリからディスクへ移し始め、これは非常に高コストになりえる。シャッフル中に最も起きやすい」と説明しています。ステージ詳細に Shuffle spill (memory) と Shuffle spill (disk) が出ていなければ、そのステージにスピルはありません。偏りについては、Summary Metrics のタスク実行時間で最大値が 75 パーセンタイルより 50% 以上大きければスキューを疑うという目安が示されています。健全なステージでは両者がほぼ同じ値になります。

読む量そのものを減らす手もあります。Delta Lake では述語プッシュダウンと列の刈り込みに加え、各データファイルの統計値を使ったファイル単位の枝刈りが働くため、JOIN の前にフィルタで絞るだけでスキャン量が大きく変わります。