python - spark 可以将数据框拆分为 topandas 的部分

Question

我有一个 1000 万条记录的数据框。我的要求是我需要在 pandas 中对这些数据进行一些操作，而且我没有内存让所有 1000 万条记录一次都在 pandas 中。所以我希望能够将它分块并在每个 block 上使用 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 部分或按记录数分成多个部分，比如说一次 100 万条。两种解决方案都可以接受，我只需要将其分成较小的 block 进行处理。

Answer 1

一种选择是将 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 block 上。这可以使用 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)