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

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

PySpark の DataFrame で CASE 式を書いてみた

こんにちは、kaorr です。

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

はじめに

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

f:id:kaorr_mob:20171205232737j: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 を出力する

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

# カラムの処理を良い感じにする col と、CASE 式用の when を import
from pyspark.sql.functions import col, when

# select で name と salary(CASE 式の結果) を指定
# when をメソッドチェーンで繋いでいく
df_result = df.select(col('name'), when(col('salary_A').isNotNull(), col('salary_A')).when(col('salary_B').isNotNull(), col('salary_B')).otherwise(0).alias('salary'))

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

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

おわりに

今回は SQL を使わずに書いてみましたが、実務レベル(?)にやるのであれば、やはり TempView を作って SQL でクエリ投げた方が良いと思います。

とはいえ、やはり Python のレイヤーでガチャガチャやるのは面白いです、はい。

f:id:kaorr_mob:20171205231110j:plain

Spark 面白いよ Spark。