TopolearnData scientists constantly cross between Spark and pandas; doing it with Arrow and the function APIs keeps it fast and scalable.
Arrow is the bridge between the JVM and Python. With spark.sql.execution.arrow.pyspark.enabled turned on (the default in Spark 4.2; older versions and some platforms need it set explicitly), df.toPandas() and spark.createDataFrame(pandas_df) transfer columnar batches instead of pickled rows, often 10 to 100 times faster. toPandas() still collects everything to the driver, so use it only on results that fit in driver memory.
For pandas logic that should run on the executors, use the pandas function APIs:
df.groupBy("store").applyInPandas(func, schema):funcreceives one group as a pandas DataFrame and returns a DataFrame. Use it for per-group models, custom resampling or anything groupwise that built-ins can't do. Each group must fit in an executor's memory.df.mapInPandas(func, schema):funcreceives an iterator of pandas DataFrames (Arrow batches) for each partition and yields DataFrames. Great for batch inference with a model loaded once per partition.df.groupBy(...).cogroup(other.groupBy(...)).applyInPandas(...)for two DataFrames.- Arrow-native variants such as
mapInArrowpass PyArrow RecordBatches.
The pandas API on Spark (import pyspark.pandas as ps) offers pandas syntax on distributed data, for example ps.read_parquet(...) and psdf.groupby(...).mean(), convenient for porting pandas code, though native Spark DataFrames give more control.
Example
from pyspark.sql import SparkSession, functions as F
spark = SparkSession.builder.master("local[2]").getOrCreate()
import pandas as pd
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
df = spark.createDataFrame(pd.DataFrame({
"store": ["A", "A", "A", "B", "B"],
"day": [1, 2, 3, 1, 2],
"sales": [10.0, 12.0, 20.0, 5.0, 7.0],
}))
def zscore(g: pd.DataFrame) -> pd.DataFrame: # runs on executors, per store
return g.assign(z=((g.sales - g.sales.mean()) / g.sales.std()).round(2))
scored = df.groupBy("store").applyInPandas(
zscore, schema="store string, day long, sales double, z double")
print(scored.orderBy("store", "day").toPandas()["z"].tolist())
# [-0.76, -0.38, 1.13, -0.71, 0.71]
def add_tax(batches): # one iterator per partition
for pdf in batches:
yield pdf.assign(gross=pdf.sales * 1.25)
print(df.mapInPandas(add_tax, "store string, day long, sales double, gross double")
.agg(F.sum("gross")).first()[0]) # 67.5Remember
- Enable Arrow for fast
toPandasandcreateDataFrame(pandas_df) applyInPandasruns pandas per group;mapInPandasper partition batchtoPandascollects to the driver: only for results that fit
Watch out for
- Calling
toPandas()on a huge DataFrame to run pandas code, instead ofapplyInPandason executors.