How to slice a pyspark dataframe in two row-wise
Spark dataframes cannot be indexed like you write. You could use head method to Create to take the n top rows. This will return a list of Row() objects and not a dataframe. So you can convert them back to dataframe and use subtract from the original dataframe to take the rest of the rows.
#Take the 100 top rows convert them to dataframe
#Also you need to provide the schema also to avoid errors
df1 = sqlContext.createDataFrame(df.head(100), df.schema)
#Take the rest of the rows
df2 = df.subtract(df1)
You can use also SparkSession instead of spark sqlContext if you work on spark 2.0+. Also if you are not interested in taking the first 100 rows and you want a random split you can use randomSplit like this:
df1,df2 = df.randomSplit([0.20, 0.80],seed=1234)
Initially I misunderstood and thought you wanted to slice the columns. If you want to select a subset of rows, one method is to create an index column using monotonically_increasing_id()
. From the docs:
The generated ID is guaranteed to be monotonically increasing and unique, but not consecutive.
You can use this ID to sort the dataframe and subset it using limit()
to ensure you get exactly the rows you want.
For example:
import pyspark.sql.functions as f
import string
# create a dummy df with 500 rows and 2 columns
N = 500
numbers = [i%26 for i in range(N)]
letters = [string.ascii_uppercase[n] for n in numbers]
df = sqlCtx.createDataFrame(
zip(numbers, letters),
('numbers', 'letters')
)
# add an index column
df = df.withColumn('index', f.monotonically_increasing_id())
# sort ascending and take first 100 rows for df1
df1 = df.sort('index').limit(100)
# sort descending and take 400 rows for df2
df2 = df.sort('index', ascending=False).limit(400)
Just to verify that this did what you wanted:
df1.count()
#100
df2.count()
#400
Also we can verify that the index column doesn't overlap:
df1.select(f.min('index').alias('min'), f.max('index').alias('max')).show()
#+---+---+
#|min|max|
#+---+---+
#| 0| 99|
#+---+---+
df2.select(f.min('index').alias('min'), f.max('index').alias('max')).show()
#+---+----------+
#|min| max|
#+---+----------+
#|100|8589934841|
#+---+----------+