はじめに

こんにちは。株式会社ジール所属の@m_akiguchiです。

普段はAWSやPythonを使った開発作業を行っています。
PySparkで条件分岐処理を実装する際、つまずいた点があったのでTipsとしてまとめます。

実行環境がない場合は、以下の記事を参考にしてみてください。
Python：Pythonによる関数の使い方
PySpark：Pythonで動かしてみるSpark入門

PySparkでの条件分岐

PySparkで条件分岐を行う場合、when、otherwiseを使用します。
基本的な書き方は以下の通りです。
when(条件,条件に合致した場合の値).otherwise(条件に合致しなかった場合の値)

例
以下の従業員テーブルを用いて年齢が30歳以上と30歳未満の社員を判断したい

従業員テーブル（t_emp）

emp_no	name	dept_no	age
001	田中太郎	005	25
002	東京花子	010	33
003	福岡次郎	001	46
004	神奈川美香	020	18
005	品川三郎	002	51

from pyspark.sql import *
from pyspark.sql.functions import *

# 上記テーブルを手で作成する場合、このように記述してください

spark = SparkSession.builder.getOrCreate()

df = spark.createDataFrame([
    ("001", 田中太郎, 005, 25),
    ("002", 東京花子, 010, 33),
    ("003", 福岡次郎, 001, 46),
    ("004", 神奈川美香, 020, 18),
    ("005", 品川三郎, 002, 51)
], ["emp_no", "name", "dept_no", "age"])

# SELECT
df = df.select(col('emp_no'),
               col('name'),
               col('age'),
               when(col('age') >= 30, lit('30歳以上')).otherwise(lit('30歳未満'))
              )

# 表示
df.show()

条件に応じて文字列や数字を出力したい場合、lit（リテラル）を使います

出力結果イメージ

emp_no	name	age	CASE WHEN age >= 30...
001	田中太郎	25	30歳未満
002	東京花子	33	30歳以上
003	福岡次郎	46	30歳以上
004	神奈川美香	18	30歳未満
005	品川三郎	51	30歳以上

カラムに別名を付けるalias（エイリアス）を使って条件分岐処理を加えたカラムに別名を付けることができます。


t_emp.select(col('emp_no'),
             col('name'),
             col('age'),
             when(col('age') >= 30,lit('30歳以上')).otherwise(lit('30歳未満')).alias('30歳以上か？')
            )

出力結果イメージ

emp_no	name	age	30歳以上か？
001	田中太郎	25	30歳未満
002	東京花子	33	30歳以上
003	福岡次郎	46	30歳以上
004	神奈川美香	18	30歳未満
005	品川三郎	51	30歳以上

#複数条件での分岐

条件分岐を行う際、複数の条件を使用したいケースがよくあると思います。
その場合、whenを複数記述することで対応できます。
when(条件1,条件に合致した場合の値).when(条件2,条件に合致した場合の値).otherwise(条件に合致しなかった場合の値)

例
以下の従業員テーブルを用いて年齢が30歳未満か50歳以上とそれ以外の社員を判断したい

従業員テーブル（t_emp）

emp_no	name	dept_no	age
001	田中太郎	005	25
002	東京花子	010	33
003	福岡次郎	001	46
004	神奈川美香	020	18
005	品川三郎	002	51

t_emp.select(col('emp_no'),
             col('name'),
             col('age'),
             when(col('age') < 30,lit('30歳未満'))
             .when(col('age') >= 50,lit('50歳以上'))
             .otherwise(lit('その他')).alias('年齢範囲')
            )

出力結果イメージ

emp_no	name	age	年齢範囲
001	田中太郎	25	30歳未満
002	東京花子	33	その他
003	福岡次郎	46	その他
004	神奈川美香	18	30歳未満
005	品川三郎	51	50歳以上

条件分岐を行う際、複数の条件を組み合わせて処理を分岐したいこともよくあると思います。
その場合、アンパサンド（&）またはバーティカルバー（|）区切りで複数条件を記述することができます。
アンパサンドを使うとAND条件、バーティカルバーを使うとOR条件となります。
when(条件1 & 条件2,条件に合致した場合の値).otherwise(条件に合致しなかった場合の値)
when(条件1 | 条件2,条件に合致した場合の値).otherwise(条件に合致しなかった場合の値)

例
以下の従業員テーブルを用いて年齢によって健康診断の内容を変更したい

従業員テーブル（t_emp）

emp_no	name	dept_no	age
001	田中太郎	005	25
002	東京花子	010	33
003	福岡次郎	001	46
004	神奈川美香	020	18
005	品川三郎	002	51

t_emp.select(col('emp_no'),
             col('name'),
             col('age'),
             when(col('age') < 35,lit('健康診断（若年層）'))
             .when((col('age') >= 35) & (col(age) < 50),lit('人間ドッグ'))
             .otherwise(lit('人間ドッグ＋メタボ健診')).alias('健康診断コース')
            )

アンパサンドやバーティカルバーを使用して複数条件を組み合わせる場合、それぞれを()で括る必要があります

出力結果イメージ

emp_no	name	age	健診コース
001	田中太郎	25	健康診断（若年層）
002	東京花子	33	健康診断（若年層）
003	福岡次郎	46	人間ドッグ
004	神奈川美香	18	健康診断（若年層）
005	品川三郎	51	人間ドッグ＋メタボ健診

#おわりに
いかがでしたでしょうか。
条件分岐が使えるようになると、PySparkを用いてできることがぐっと広がります。
複数条件の書き方やlit()を使用する場合など、実際書いてみるとつまずきやすいポイントもまとめたため、
役立てていただけると嬉しいです。