はじめに
こんにちは。株式会社ジール所属の@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()を使用する場合など、実際書いてみるとつまずきやすいポイントもまとめたため、
役立てていただけると嬉しいです。