Pysparkで、同じ値だけを各行に入れた列を新設したい

今回は少しマイナーなTipsかもしれませんが、意外と「あれどうやるんだっけ？」となることが多いので書いてみます。

今回も例によってサンプルデータを用意します。

sample_data = [{'name': 'Alice',"age": '2',}, 
               {'name': 'Bob',"age": '50'},
               {'name': 'Matilda',"age": '12'}]

df = spark.createDataFrame(sample_data)
df.show()

【結果】
+---+-------+
|age|   name|
+---+-------+
|  2|  Alice|
| 50|    Bob|
| 12|Matilda|
+---+-------+

ここに、birthplace（出身地）という列を新設しようと思います。
値となる地名は何でもいいですが、せっかく名前にマチルダを追加したのでここでは映画『レオン』に則って「NewYork」としましょうか。
3人とも、全員ニューヨーク出身として、NewYorkという値が各行に入った列を追加します。

litで値を囲むのがミソ

結論、こんな感じで書くことが出来ます！

from pyspark.sql.functions import lit
df_add_birthplace = (
    df
    .withColumn("birthplace",lit("NewYork"))
)
df_add_birthplace.show()

【結果】
+---+-------+----------+
|age|   name|birthplace|
+---+-------+----------+
|  2|  Alice|   NewYork|
| 50|    Bob|   NewYork|
| 12|Matilda|   NewYork|
+---+-------+----------+

列を追加するコマンドとして.withColumnがある、というのは割と有名ですよね。
ここでは

df.withColumn("新設するカラム名",lit("新設列に入れる値"))

といった文法になります。

想定されるユースケース：名前からメアドを生成

といっても、Pysparkで行うETL作業の中で、元データから必要なものだけに削ぐことはあっても追加する、ってことは経験上あまりないのかなと思っています。

そのなかでも、想定されるユースケースとしては、例えば「メールアドレスの生成」があるかなと思います！

sample_data = [{'name': 'Alice',"age": '2',}, 
               {'name': 'Bob',"age": '50'},
               {'name': 'Matilda',"age": '12'}]
df = spark.createDataFrame(sample_data)

#「mail_domain」列を追加する
from pyspark.sql.functions import lit
df_add_mail_domain = (
    df
    .withColumn("mail_domain",lit("@amail.com"))
)
df_add_mail_domain.show()

【結果】
+---+-------+-----------+
|age|   name|mail_domain|
+---+-------+-----------+
|  2|  Alice| @amail.com|
| 50|    Bob| @amail.com|
| 12|Matilda| @amail.com|
+---+-------+-----------+

そして、nameとmail_domainをconcatで結合すると、晴れてメールアドレスの出来上がりです！
（今回は架空ドメインの「@amail.com」を使うことにします）

from pyspark.sql.functions import col,concat
df_add_mail_address = (
    df_add_mail_domain
    .withColumn("mail_address",concat(col("name"),col("mail_domain"))
    )
)
df_add_mail_address.show()

【結果】
+---+-------+-----------+-----------------+
|age|   name|mail_domain|     mail_address|
+---+-------+-----------+-----------------+
|  2|  Alice| @amail.com|  Alice@amail.com|
| 50|    Bob| @amail.com|    Bob@amail.com|
| 12|Matilda| @amail.com|Matilda@amail.com|
+---+-------+-----------+-----------------+

ちなみに文法はこんな感じ。

(データフレーム).withColumn("新設するカラム名",concat(col("結合列①"),col("結合列②"))

いかがでしょうか？
この他にも「この列の値に特定の値を付与したいよね」というニーズには、上記のような「litで新列作る」→「concatで結合」の流れを踏むことで応えられる気がします。

あ、マチルダにメールは送らないでくださいね笑

このシリーズの目的

クラウドストレージに蓄積された生データ（CSV,JSON等）を加工するのに必要なPysparkの知識を溜めていく、まさに「Pysparkに関するトリビアのDalta Lake」を目指しています。
Pysparkに関する記事はまだ海外のものが多く、このシリーズが私のような国内の駆け出しデータエンジニアの一助になることを願っています。