Pandas データフレームを Spark データフレームへスキーマを指定して変換する際にはカラム順に注意

概要

Pandas データフレーム（以下、Pandas）を Spark データフレーム（以下、Spark df）へスキーマを指定して変換する際には、指定したスキーマ順で変換する仕様であり、変換前の Pandas とカラム名が変わってしまう場合あることに注意が必要です。

下記図にて、上の結果が Pandas の結果であり、下の結果が Spark df の結果です。Pandas の geder と profession が Spark df に変換すると値はそのままであるがカラム名が反対になっていることが確認できます。これは、schema にて profession を gender より先に指定していることが要因のようです。

import pandas as pd

data = {
    "name": ["John", "Jane", "Sam", "Peter"],
    "age": [25, 32, 18, 47],
    "gender": ["Male", "Female", "Male", "Male"],
    "profession": ["Engineer", "Doctor", "Student", "Lawyer"],
}

pandas_df = pd.DataFrame(data)
pandas_df.display()

schema = """
name string,
age int,
profession string,
gender string
"""

df = spark.createDataFrame(pandas_df,schema)
df.display()

スキーマを指定せずに Pandas から Spark df に変換も可能ですが、データがない場合に次のようなエラーが発生するなど想定外のエラーが発生することがありスキーマ指定が必要な場合があります。

ValueError: can not infer schema from empty dataset

本記事ではカラムを想定通りに取得する方法として次の方法を紹介します。私は、プログラムでカラム順の保証ができたため、 1 の方法を採用しました。

1. Pandas のカラム順に揃えたスキーマを指定する方法
2. Pandas と スキーマのカラム順を並び替えて変換後に想定のカラム順に戻す方法
3. Pandas を辞書型変数に変換して Spark データフレームに変換する方法

対応方法

事前準備

import pandas as pd

data = {
    "name": ["John", "Jane", "Sam", "Peter"],
    "age": [25, 32, 18, 47],
    "gender": ["Male", "Female", "Male", "Male"],
    "profession": ["Engineer", "Doctor", "Student", "Lawyer"],
}

pandas_df = pd.DataFrame(data)
pandas_df.display()

1. Pandas のカラム順に揃えたスキーマを指定する方法

schema = """
name string,
age int,
gender string,
profession string
"""

df = spark.createDataFrame(pandas_df,schema)
df.display()

2. Pandas と スキーマのカラム順を並び替えて返還後に想定のカラム順に戻す方法

# pandas_df が変更されないように Pandas をコピー
pandas_df_02 = pandas_df.copy()

# 並び替え前のカラム順のリストを取得
cols = list(pandas_df_02.columns).copy()

# Pandas のカラム順を昇順に変更
pandas_df_02 = pandas_df_02.reindex(sorted(pandas_df_02.columns), axis=1)
# pandas_df_02.display()

# 昇順に並び替えた Spark データフレームのスキーマを取得
schema = """
name string,
age int,
gender string,
profession string
"""
tmp_df = spark.createDataFrame([],schema)
tmp_df = tmp_df.select(sorted(tmp_df.columns))
sorted_schema = tmp_df.schema
# print(sorted_schema)

# 並び替え済みのデータにより Spark データフレームに変換し、カラム順を修正
df = spark.createDataFrame(pandas_df_02,sorted_schema)
df = df.select(cols)
df.display()

3. Pandas を辞書型変数に変換して Spark データフレームに変換する方法

schema = """
name string,
age int,
profession string,
gender string
"""

df = spark.createDataFrame(pandas_df.to_dict(orient='records'),schema)
df.display()