切片Spark的DataFrame SQL的行(pyspark)

Question

我有一个Spark的Dataframe parquet文件，可以被Spark读取如下

df = sqlContext.read.parquet('path_to/example.parquet')
df.registerTempTable('temp_table')

我想切片我的数据框，df，按行(即相当于df.iloc[0:4000], df.iloc[4000:8000]等熊猫数据框)，因为我想把每个小块熊猫数据框以后工作。我只知道如何使用sample随机分数，即

df_sample = df.sample(False, fraction=0.1) # sample 10 % of my data
df_pandas = df_sample.toPandas()

我将是伟大的，如果有一种方法切片我的数据框df行。

Answer 1

您可以使用单调递增_id()添加一个ID列到您的数据框，并使用它来获得任何大小的工作集

import pyspark.sql.functions as f
# add an index column
df = df.withColumn('id', f.monotonically_increasing_id())
# Sort by index and get first 4000 rows
working_set = df.sort('id').limit(4000)

然后，您可以使用subtract()从数据框架中删除工作集。

# Remove the working set, and use this `df` to get the next working set
df = df.subtract(working_set)

冲洗并重复，直到处理完所有行。这不是理想的做事方式，但确实有效。考虑过滤掉要在Pandas中使用的Spark数据框架。