PySpark's "DataFrameLike" type vs pandas.DataFrame

user_3945382 0
Viewed 595

Spark 3.1 introduced type hints for python (hooray!) but I am puzzled as to why the return type of the toPandas method is "DataFrameLike" instead of pandas.DataFrame - see here: https://github.com/apache/spark/blob/master/python/pyspark/sql/pandas/conversion.pyi

Because of this mypy throws all sorts of errors if I try to use any of the pandas df methods on an object that's the result of calling toPandas. For example

df = spark_df.toPandas()
df.to_csv(out_path, index=False)

results in the error message

error: "DataFrameLike" has no attribute "to_csv"

What's going on here?

pyspark apache-spark-sql apache-spark type-hinting python
2 Answers
user_135852690

To fix the mypy warnings:

cast has no effect at runtime, but it tells mypy to treat it as a real pandas.DataFrame for the purposes of type-checking.

I like the other answers here, and maybe you can fix it without this cast trick/hack, but I'm giving this as another option

import pandas as pd
from typing import cast

df = cast(pd.DataFrame, spark_df.toPandas())
df.to_csv(out_path, index=False)
user_1460410
Related