PySparkでDataFrameのカラム名をrename

概要

PySparkでSELECTしたカラムの名前を変更する方法を確認する。

バージョン情報

• Spark 2.4.3

サンプルデータ

カラムfoo, bar, bazの3つを持つtmp_exampleテーブルを用意。

import pandas as pd
import numpy as np
np.random.seed(0)

# 0〜100までの値がランダムに入った10行のデータを用意
pandas_df = pd.DataFrame(
    np.random.randint(0, 100, (10, 3)),
    columns=['foo', 'bar', 'baz'])

# SparkのDataFrameに変換してtmp_exampleに保存
df = spark.createDataFrame(pandas_df)
df.registerTempTable('tmp_example')

このテーブルを参照して、fooというカラムをhogeという名前に変更する例を記述する。

sqlで書く

慣れないうちはSQLですべて解決してしまうと楽。

query = 'select foo as hoge, bar, baz from tmp_example'
spark.sql(query).take(3)
    #=> [Row(hoge=44, bar=47, baz=64),
    #=>  Row(hoge=67, bar=67, baz=9),
    #=>  Row(hoge=83, bar=21, baz=36)]

selectExpr

selectExprを使うとselectする際にSQLの表現が使える。

例えば下記はasでfooをhogeに書き換えている。

spark.table('tmp_example') \
     .selectExpr('foo as hoge', 'bar', 'baz') \
     .take(3)
    #=> [Row(hoge=44, bar=47, baz=64),
    #=>  Row(hoge=67, bar=67, baz=9),
    #=>  Row(hoge=83, bar=21, baz=36)]

計算結果を取ったりUDFを呼び出したりもできる。

spark.table('tmp_example') \
      .selectExpr('foo * 2 as hoge', 'sqrt(bar) as fuga') \
      .take(3)
    #=> [Row(hoge=88, fuga=6.855654600401044),
    #=>  Row(hoge=134, fuga=8.18535277187245),
    #=>  Row(hoge=166, fuga=4.58257569495584)]

alias

colのaliasを使って名前を変える。

from pyspark.sql import functions as F
spark.table('tmp_example') \
      .select(F.col('foo').alias('hoge'), 'bar', 'baz') \
      .take(3)
    #=> [Row(hoge=44, bar=47, baz=64),
    #=>  Row(hoge=67, bar=67, baz=9),
    #=>  Row(hoge=83, bar=21, baz=36)]

toDFでまとめて変える

column_names = ['hoge', 'fuga', 'piyo']
spark.table('tmp_example') \
      .toDF(*column_names) \
      .take(3)
    #=> [Row(hoge=44, fuga=47, piyo=64),
    #=>  Row(hoge=67, fuga=67, piyo=9),
    #=>  Row(hoge=83, fuga=21, piyo=36)]

withColumnRenamed

withColumnRenamedで変更前、変更後のカラム名を指定してrenameできる。

spark.table('tmp_example') \
      .withColumnRenamed('foo', 'hoge').take(3)
    #=>  [Row(hoge=44, bar=47, baz=64),
    #=>   Row(hoge=67, bar=67, baz=9),
    #=>   Row(hoge=83, bar=21, baz=36)]

カラム名を一括で変換する

columnsでカラム名が取れる。

df = spark.table('tmp_example')
df.columns
    #=>  ['foo', 'bar', 'baz']

すべてのカラム名の末尾に _1 を付けてみる。

new_columns = ['%s_1' % c for c in df.columns] 

df.toDF(*new_columns).take(3)
    #=> [Row(foo_1=44, bar_1=47, baz_1=64),
    #=>  Row(foo_1=67, bar_1=67, baz_1=9),
    #=>  Row(foo_1=83, bar_1=21, baz_1=36)]

改定履歴

Author: Masato Watanabe, Date: 2019-07-13, 記事投稿