Randomly access chunk of data from pyspark dataframe

user_12354392 0
Viewed 83

I am trying to implement pagination using the pyspark dataframe. Earlier I was thinking of having page numbers calculated and then slice the chunk of rows based on page number and return.

but after understanding a little bit about pyspark I am not able to get this working as pyspark does not allow accessing rows from middles/ randomly.

I am new to pyspark, what I am trying to implement here looks like.

I have 5000 rows in the pyspark dataframe. I want to have 100 rows per page and fetch only 10 pages at a time so that it does not impact local memory.

pyspark python-3.x pagination
1 Answers

I can think of two ways we can solve this problem, both involve row_number(). We can use row_number() over a dummy window to create a monotonically increasing integer value for each row.

With method 1 we can use our row number to filter the pages we want with a query on the backend. Credit for this method here https://stackoverflow.com/a/67488679/19167644

Method 1

import pyspark.sql.functions as F
from pyspark.sql.window import Window
w = Window().orderBy(F.lit('A'))
df = df.withColumn('row_num', F.row_number().over(w))
offset, limit = 10, 9
data = df.where(F.col('row_num').between(offset, offset + limit))

With method 2 we can create a column that groups the rows in groups of 10, and query that column directly.

Method 2

import pyspark.sql.functions as F
from pyspark.sql.window import Window
w = Window().orderBy(F.lit('A'))
df = df.withColumn('row_num', F.floor(F.row_number().over(w) / 10) )

Edit: notes on

w = Window().orderBy(F.lit('A'))

For row_number() to generate a new integer for each row, we have to supply a partition that has no relation to our data. This line generates a nonsensical or dummy window for row_number() to reference. Since there is no relation, each row is its own window, and thus gets a new integer.

user_191676440
Related