SQL で書いても DataFrame で書いても同じ計画になり、ウィンドウ関数は行を減らさない

一言でいうと

SQL文字列とDataFrame APIは、同じエンジンが同じプランに変換する2つの入り口です。どちらが速いかではなく、どちらが読みやすいかで選びます。集計は行をまとめて減らしますが、ウィンドウ関数は行をそのままにして隣の行を見ます。そして、ウィンドウの結果はフレームをどう設定したかによって決まります。

なぜ入り口が2つあるのか

データチームには、SQLで考える人とPythonで考える人が混在しています。すると、すぐに「SQLで書いたほうが速いらしい」「DataFrameのほうが最適化がよく効くらしい」といった話が出回ります。どちらも間違いです。

Spark SQLガイドの最初の段落が答えです。Spark SQLはRDDよりデータと計算の構造を多く把握しており、その情報でより多くの最適化を行います。そして結果を計算するときは、どのAPIや言語で表現しても同じ実行エンジンを使います。ドキュメントは、この統一のおかげで、トランスフォーメーションごとに自然なAPIを選んで行き来できると書いています。

同じエンジンという言葉は具体的です。SQL文字列はパースされて論理プランになり、DataFrameのメソッド呼び出しもそのまま論理プランを積み上げます。そこからは、同じアナライザー・同じオプティマイザー・同じ物理プランナーを通ります。そのため、同じ質問を2通りの方法で書いてexplain()を出力すると、物理プランは同じです。このラボでそれを自分で比べます。

どう動くのか

SQLでDataFrameを呼び出すには、名前が必要です。入門ドキュメントは、createOrReplaceTempViewでDataFrameを一時ビューとして登録し、spark.sqlでクエリする形を示しています。同じ箇所には、一時ビューはセッションスコープなので、作成したセッションが終了すると消えると書かれています。ビューはデータをコピーしません。プランに付けた名前にすぎません。

orders.createOrReplaceTempView("orders")
by_sql = spark.sql("""
  SELECT date_trunc('month', order_ts) AS month, sum(qty) AS units
  FROM orders WHERE status = 'paid' GROUP BY 1""")
by_api = (orders.where(F.col("status") == "paid")
          .groupBy(F.date_trunc("month", "order_ts").alias("month"))
          .agg(F.sum("qty").alias("units")))

集計の物理プランは、ほぼいつも同じ形です。EXPLAINドキュメントの例を見ると、下から順にHashAggregate(... partial_sum ...) → Exchange hashpartitioning(k, 200) → HashAggregate(... sum ...)の順です。各パーティションがまず自分の担当分を部分合計に減らし、キーでシャッフルしたあと、集まった部分合計を再び足します。シャッフルで送るのは元の行ではなく、キーごとに1行に減らした部分合計です。そのため、合計・件数のような集計は、データが大きくてもシャッフルが思ったより小さくなります。例の200はシャッフルパーティション数で、シャッフルを扱うモジュールで別に見ます。

ウィンドウ関数(行を減らさない集計)

groupByはカテゴリごとに1行を残します。「カテゴリ別の上位3商品」は、そのやり方では解けません。商品の行はそのままにして、同じカテゴリの中で何位かを隣のカラムに付ける必要があります。これがウィンドウ関数です。SQLリファレンスのウィンドウ関数のページは、ランキング関数(RANK・DENSE_RANK・ROW_NUMBERなど)と分析関数(LAG・LEAD・FIRST_VALUEなど)を区別し、フレームをROWSまたはRANGEで記述する構文を定義しています。

ランキング関数の違いは、同点のときに出ます。rankのドキュメントの説明のとおり、dense_rankは同点のあとに空きの順位を残さず、rankは3つが同率2位なら、その次を5位とします。「上位3」をrank <= 3で絞り込むと、同点のために4行以上が出ることがあり、row_numberで絞り込むと、同点のうち誰が入るかが並び順に委ねられます。どちらが正しいかは、業務が決めます。

もう1つの静かな落とし穴は、デフォルトのフレームです。Windowのドキュメントは、並び順がなければパーティション全体(行基準で、最初から最後まで)を、並び順があれば範囲基準で最初から現在の行まで伸びるフレームを、デフォルトとして使うと書いています。rangeBetweenのドキュメントによると、範囲の境界は行の位置ではなくORDER BYの値を基準にします。そのため、日付で並べて累積合計を出すと、同じ日付の行同士が互いを「現在の行まで」に含めてしまい、同じ累積値を一度に受け取ります。1行ずつ増える累積合計が欲しければ、rowsBetween(Window.unboundedPreceding, Window.currentRow)で行基準のフレームを書く必要があります。

正確な数と近似した数

ユニークな顧客数を正確に数えるには、すべての顧客IDを集めて重複を除く必要があるため、値そのものをシャッフルしなければなりません。approx_count_distinctは許容相対標準偏差rsdを受け取って推定し、デフォルトは0.05です。ドキュメントは、rsdが0.01より小さくなければならないなら、むしろcount_distinctのほうが効率的だと書いています。組み込み集計関数の一覧は、この推定がHyperLogLog++で行われると明かしています。ダッシュボードの日次訪問者のように5%の誤差が許容できる場所に使い、精算のように1人でも間違ってはいけない場所には使いません。

現場での姿

SQL文字列のエラーは、実行時に出ます。カラム名を間違えても、Pythonのエディターは何も言わず、spark.sqlを呼んだ瞬間に分析エラーが出ます。DataFrame APIも、カラムを文字列で書けば同じです。どちらの場合も、アクションの前の分析の段階で引っかかるので、テストは小さなデータで一度最後まで実行するのが最も安上がりです。

2つの方法を混ぜるのが普通です。複雑なウィンドウやテーブルの結合はSQLで、カラムを繰り返し作る処理はPythonのループとDataFrame APIで書く、といった具合です。プランが同じなので、混ぜても損はありません。

一時ビューはセッションの外からは見えません。別のアプリや別のセッションから同じ名前で呼ぶと、テーブルが見つかりません。アプリ間でデータを渡すには、ファイルやテーブルとして書き出す必要があります。

実務で本当に大切なこと

SQLとDataFrameは、同じエンジン、同じプランです。性能ではなく読みやすさで選び、疑わしければexplainで比べます。
集計は、部分集計 → シャッフル → 最終集計です。シャッフルされるのは元の行ではなく、減らした部分結果です。
ウィンドウは行を減らしません。上位Nはランキング関数で順位を付け、同点のルールを業務に合わせて選びます。
累積合計はフレームを明示します。並び順があるウィンドウのデフォルトは範囲フレームなので、同じ値の行が一度にまとめられます。
近似ユニーク数は、誤差が許容できる場所にだけ使います。デフォルトのrsdは0.05です。

次のラボですること

注文データを一時ビューとして登録してSQLで月別売上を出し、同じ結果をDataFrame APIで書き直したあと、両方のプランが同じであることをexplainで確認します。ウィンドウ関数でカテゴリ別の上位3商品を取り出し、行基準のフレームで累積合計を、lagで前月比の変化を計算します。最後に、ユニーク顧客数を近似関数と正確な方法でそれぞれ数えて誤差を測り、レポートにまとめます。