モブプロな人たちのブログ

横浜で Web サービス開発しているエンジニアの日記です。Python 大好き Flask 大好き。たまに SpriteKit でゲーム開発も。

PySpark の DataFrame で CASE 式を書いてみた (SQL 編)

こんにちは、kaorr です。

一人アドベントカレンダー6日目です。

はじめに

先日構築した PySpark の検証環境を使って、DataFrame の CASE 式を書いてみようと思います。(その2)

今回は、SQL を使います。

f:id:kaorr_mob:20171206165539j:plain

サンプルデータの作成と要件

前回とほとんど同じですが、こんな感じになります。

# createDataFrame でデータフレームを作成する
df = spark.createDataFrame([('yamada', None, 200), ('tanaka', 300, None), ('takeda', 400, 200), ('suzuki', None, None)], ['name', 'salary_A', 'salary_B'])

# データフレームの中身を見る
df.show()

# ↓こんな感じの出力があれば OK
# 
# +------+--------+--------+
# |  name|salary_A|salary_B|
# +------+--------+--------+
# |yamada|    null|     200|
# |tanaka|     300|    null|
# |takeda|     400|     200|
# |suzuki|    null|    null|
# +------+--------+--------+
  • 社員ごとの給与を出力したい

  • 給与パターン A と給与パターン B が存在する

  • A がある場合は A を出力する (B は無視)

  • A がなくて B がある場合は B を出力する

  • 両方ない場合は 0 を出力する

  • [追加] カラムは増減する可能性がある (salary_C が追加されたり、salary_B が削除されたりする)

CASE 式の適用とデータの確認

# salary_ から始まるカラムを取得しておく
salary_columns = [column for column in df.columns if column.startswith('salary_')]

# salary_columns の中身はこんな感じ
# ['salary_A', 'salary_B']

# データフレームから temp_table を作成
df.createTempView('temp_table')

# CASE 式の WHEN 〜 THEN 〜 の部分を配列で生成
when_list = ['when `{0}` is not null then `{0}`'.format(column) for column in salary_columns]

# when_list の中身はこんな感じ
# ['when `salary_A` is not null then `salary_A`', 'when `salary_B` is not null then `salary_B`']

# SQL 文の全体を生成
sql_text = 'select name, case {0} else 0 end as salary from temp_table'.format(' '.join(when_list))

# sql_text の中身はこんな感じ
# 'select name, case when `salary_A` is not null then `salary_A` when `salary_B` is not null then `salary_B` else 0 end as salary from temp_table'

# クエリ実行!!
df_result = spark.sql(sql_text)

# df_result の中身を見る
df_result.show()

# ↓こんな感じの出力があれば OK
# 
# +------+------+
# |  name|salary|
# +------+------+
# |yamada|   200|
# |tanaka|   300|
# |takeda|   400|
# |suzuki|     0|
# +------+------+

おわりに

前回と比べると、やはり SQL で書いた方がしっくりきますね。

また、今回はカラムが増減する可能性がある場合でも、汎用的に使えるように書いてみました。

SQL 文を Python でガチャガチャいじりながら生成するのは、面白いですね。

f:id:kaorr_mob:20171206170018j:plain

Python いいよ Python。