可能将数据框拆分为topandas的火花

时间:2018-10-26 15:52:22

标签: python pandas apache-spark

我有一个1000万个记录数据框。我的要求是我需要对熊猫中的这些数据进行一些操作,而我没有足够的内存来一次存储所有1000万条记录。所以我希望能够对它进行分块并在每个块上使用toPandas

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

如何将我的数据框分成相等的x部分或按记录计数分成几部分,一次说一百万。两种解决方案都是可以接受的,我只需要对它进行较小的处理即可。

1 个答案:

答案 0 :(得分:3)

一种选择是将toLocalIteratorrepartitionmapPartitions结合使用。

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

通过使用toLocalIterator,一次仅将一个分区收集到驱动程序。

我认为更可取的另一种选择是,将您的工作分配到每个分区中大熊猫块上的整个群集中。可以使用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)