pysparkでリテラル列を追加(pyspark)

pysparkでリテラル列を追加したかったのでメモ。

確認環境

Welcome to
      ____              __
     / __/__  ___ _____/ /__
    _\ \/ _ \/ _ `/ __/  '_/
   /__ / .__/\_,_/_/ /_/\_\   version 2.4.1
      /_/

Using Python version 2.7.5 (default, Apr  9 2019 14:30:50)
SparkSession available as 'spark'.

コード

# 読み込むファイル
[hoge@localhost spark_script]$ cat ../names.csv
John
Mary
Taylor
Smith

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from pyspark import SparkConf
from pyspark import SparkContext
from pyspark.sql import SparkSession
from pyspark.sql.types import *
from pyspark.sql.functions import lit


def main():
    conf = SparkConf()
    conf.setAppName('example')
    sc = SparkContext(conf=conf)
    spark = SparkSession(sc)

    struct = StructType([StructField('name', StringType(), False)])

    df = spark.read.csv('/home/hoge/names.csv', schema=struct)
    
    # リテラル値の列を追加する
    # withColumnは列を追加した新しいDataFrameを返す
    df_age_added = df.withColumn("age", lit(20))
    df_age_added.show()


if __name__ == '__main__':
    main()

実行結果

[hoge@localhost spark_script]$ spark-submit literal_test.py
...
+------+---+
|  name|age|
+------+---+
|  John| 20|
|  Mary| 20|
|Taylor| 20|
| Smith| 20|
+------+---+
...

参考にしました