我有一个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部分或按记录计数分成几部分,一次说一百万。两种解决方案都是可以接受的,我只需要对它进行较小的处理即可。
答案 0 :(得分:3)
一种选择是将toLocalIterator
与repartition
和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
通过使用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)