sql - 按行切片 Spark 的 DataFrame SQL (pyspark)

Question

我有一个 Spark 的 Dataframe parquet 文件，可以通过 spark 读取，如下所示

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

我想切片我的数据帧， df ，按行(即相当于 Pandas 数据帧中的 df.iloc[0:4000], df.iloc[4000:8000] 等)，因为我想将每个小块转换为 Pandas 数据帧，以便稍后处理。我只知道如何使用 sample随机分数，即

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

如果有一种方法可以对我的数据帧进行切片，我会很棒 df按行。提前致谢。

Answer 1

您可以使用 monotonically_increasing_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 数据框。