Spark の Dataframe を１レコードのみ Python の辞書型変数に変換する際の注意事項

概要

SparkのDataframeをPythonの辞書型変数に変換する際、最初の1レコードのみを対象とする場合があります。このとき、Dataframeが空である可能性を考慮することが重要です。以下に、collectメソッドの出力結果の最初のインデックスの値を取得し、asDictメソッドで変換するプログラムを示します。

data = [
    ("Alice", 25),
]
df = spark.createDataFrame(data, ["Name", "Age"])

df_data = df.collect()
df_data[0].asDict()

Dataframeが空の場合、以下のようなエラーが発生します。これは、collectメソッドの出力結果がリスト型変数であり、存在しないインデックスを指定したためです。

# 空のDataframeを作成
empty_df = df.filter("1=0")
empty_df_data = empty_df.collect()

empty_df_data[0].asDict()

IndexError: list index out of range

通常、このようなコードを記述することは少ないですが、備忘録としてその対応方法を本記事にまとめておきます。

コードと実行結果

事前準備

data = [
    ("Alice", 25),
]
df = spark.createDataFrame(data, ["Name", "Age"])

基本的な動作確認

df_data = df.collect()
df_data[0].asDict()

{'Name': 'Alice', 'Age': 25}

エラーケースのコードを実行

以下のコードを実行すると、empty_df_data[0].asDict()の部分でエラーが発生します。これは、empty_df_dataが空のリストであり、empty_df_data[0]が存在しないためです。

# 空のDataframeを作成
empty_df = df.filter("1=0")
empty_df_data = empty_df.collect()

empty_df_data[0].asDict()

IndexError: list index out of range

データフレームが空であることを考慮した改善

以下のコードでは、empty_df.limit(1).count() > 0によりDataframeが空でないことを確認しています。Dataframeが空であれば、"empty_df_data is empty"と表示され、エラーは発生しません。

empty_df = df.filter("1=0")

if empty_df.limit(1).count() > 0:
    empty_df_data = empty_df.collect()
    empty_df_data[0].asDict()
else:
    print("empty_df_data is empty")