试试这个：

from pyspark.sql import functions as F
 sample_dict = {}

 # Read the parquet file
 df = spark.read.parquet("parquet file")

 # add the partition_number as a column
 df = df.withColumn('partition_num', F.spark_partition_id())
 df.persist()

 total_partition = [int(row.partition_num) for row in 
 df.select('partition_num').distinct().collect()]

 for each_df in total_partition:
     sample_dict[each_df] = df.where(df.partition_num == each_df)

字符串

我不相信Spark让你抵消或分页您的数据。
但你可以添加一个索引，然后分页，首先：

from pyspark.sql.functions import lit
    data_df = spark.read.parquet(PARQUET_FILE)
    count = data_df.count()
    chunk_size = 10000

    # Just adding a column for the ids
    df_new_schema = data_df.withColumn('pres_id', lit(1))
    
    # Adding the ids to the rdd 
    rdd_with_index = data_df.rdd.zipWithIndex().map(lambda (row,rowId): (list(row) + [rowId+1]))
    
    # Creating a dataframe with index
    df_with_index = spark.createDataFrame(rdd_with_index,schema=df_new_schema.schema)
    
    # Iterating into the chunks
    for page_num in range(0, count+1, chunk_size):
        initial_page = page_num*chunk_size
        final_page = initial_page + chunk_size 
        where_query = ('pres_id > {0} and pres_id <= {1}').format(initial_page,final_page)
        chunk_df = df_with_index.where(where_query).toPandas()
        train_on_batch(chunk_df) # <== Your function here

字符串
这不是最佳的，因为使用了pandas Dataframe ，它会严重利用spark，但会解决你的问题。
如果id影响到您的功能，请不要忘记删除它。