PySpark (Databricks) で使える3種類の DataFrame を相互変換する


Databricks(PySpark) を使っていると、以下の3種類の DataFrame が使える。

  • pandas.DataFrame : 皆様ご存知いつもの DataFrame
  • pyspark.sql.DataFrame : PySpark の DataFrame と言えばこれだった
  • pyspark.pandas.DataFrame : PySpark 3.2 からこれも追加された

pandas.DataFrame は、お馴染みの DataFrame なので普通に使えるが、
もちろん Spark の恩恵は受けられない。

pyspark.sql.DataFrame は、pandas.DataFrame とは API が違うのだが、
まぁ同じようなことも出来る。parquet 形式ファイルから(いつものように)データを読み出すと、普通はこの形式になるので、やはり使うことは一番多いのだろう。

pyspark.sql.DataFrame と言えば、createOrReplaceTempView() で SparkSession に貼り付けて、pyspark.sql(query) でクエリ発行して、その結果をまた pyspark.sql.DataFrame として受け取って、という処理が非常にやりやすくて泣ける。

ただ、やっぱり pandas.DataFrame みたいな処理もしたいよ〜と思う時もあり、そして逆に実際に pandas.DataFrame を使ってると Spark の恩恵も受けたいよね〜と思う時もあり。

そんな時は pyspark.pandas.DataFrame ですよ。

まぁ完璧に pandas.DataFrame じゃないんだけど、ほぼ Spark で動く pandas.DataFrame だと思って良い。

このように便利になってきているのだが、3種類も DataFrame があるとちょっとややこしいし、アレをコレに変換したいなんて時に「あれ?どうやるんだっけ?」とイソイソとググることになる。

なので、備忘録として3種類の変換方法をまとめておく。

pyspark.sql.DataFrame から変換

pyspark.sql.DataFrame → pandas.DataFrame

df_pandas = df_sparksql.toPandas()

pyspark.sql.DataFrame → pyspark.pandas.DataFrame

df_sparkpandas = df_sparksql.to_pandas_on_spark()

pandas.DataFrame から変換

pandas.DataFrame → pyspark.sql.DataFrame

spark は SparkSession のインスタンス

spark.conf.set('spark.sql.execution.arrow.enabled', 'true')
df_sparksql = spark.createDataFrame(df_pandas)

pandas.DataFrame → pyspark.pandas.DataFrame

import pyspark.pandas as ps

df_sparkpandas = ps.from_pandas(df_pandas)

pyspark.pandas.DataFrame から変換

pyspark.pandas.DataFrame → pyspark.sql.DataFrame

df_sparsql = df_sparkpandas.to_spark()

pyspark.pandas.DataFrame → pandas.DataFrame

df_pandas = df_sparkpandas.to_pandas()

toPandas() と似ていて混乱する!

ちなみに parquet 形式ファイルからも直接 pyspark.pandas.DataFrame に読み込めるので、必ず pyspark.sql.DataFrame を経由するという訳でもない。