spark possible to split dataframe into parts for topandas

Question

I have a 10million record dataframe. My requirement is I need to do some operations on this data in pandas, and I do not have the memory for all 10million records to be in pandas at once. So I want to be able to chunk it and use toPandas on each chunk

df = sqlContext.sql("select * from db.table")
#do chunking to take X records at a time
#how do I generated chunked_df?
p_df = chunked_df.toPandas()
#do things to p_df

How do I chunk my dataframe into either equal x-parts or into parts by record count, say 1 million at a time. Either solution is acceptable, I just need to process it in smaller chunks.

Answer 1

One option is to use toLocalIterator in conjunction with repartition and mapPartitions.

import pandas as pd

columns = spark_df.schema.fieldNames()
chunks = spark_df.repartition(num_chunks).rdd.mapPartitions(lambda iterator: [pd.DataFrame(list(iterator), columns=columns)]).toLocalIterator()
for pdf in chunks:
    # do work locally on chunk as pandas df

By using toLocalIterator, only one partition at a time is collected to the driver.

Another option, which in my opinion is preferable, is to distribute your work across the cluster on the pandas chunks in each partition. This can be achieved using pandas_udf:

from pyspark.sql.functions import spark_partition_id, pandas_udf, PandasUDFType

@pandas_udf(result_schema, PandasUDFType.GROUPED_MAP)
def transform_pandas_df_chunk(pdf):
    result_pdf = ...
    # do ditributed work on a chunk of the original spark dataframe as a pandas dataframe
    return result_pdf

spark_df = spark_df.repartition(num_chunks).groupby(spark_partition_id()).apply(transform_pandas_df_chunk)