Extract First and last N rows from PySpark DataFrame

In PySpark, extracting the first or last N rows from a DataFrame is a common requirement in data analysis and ETL pipelines. PySpark provides multiple ways to achieve this, either by using built-in DataFrame functions like limit(), head(), and tail(), or by applying window functions with row_number() when working with grouped or partitioned data.

The approach you choose depends on whether you need rows from the entire DataFrame or from each group within the DataFrame.

In this article, we’ll explore:

What are the different ways to extract the top N rows
How to extract the last N rows
Extracting the top N rows per group using window functions
Extracting the last N rows per group using window functions
Key points and FAQs

Key Points-

Quick inspection – Extracting the first few rows helps in understanding the schema and structure of large datasets.
head(n) – Returns the first n rows as a Python list of Row objects.
first() – Returns only the very first row as a single Row object, similar to head(1) but without the list wrapper.
limit(n).collect() – Limits the DataFrame to n rows and collects them into a Python list.
tail(n) – Fetches the last n rows directly, useful for checking recent or end records.
collect()[-n:] – Retrieves the whole DataFrame as a list and slices the last n rows, though less efficient for large datasets.
Grouped extraction – Use Window functions (row_number, rank, dense_rank) with partitionBy() to get the first or last N rows per group.
Ordering matters – Use orderBy(desc()) for top rows and orderBy(asc()) for bottom rows when working with groups.
Scalability – Window-based approaches scale well across distributed Spark clusters compared to Python slicing.
Use cases – Commonly used for sampling, identifying top performers, retrieving recent records, debugging ETL pipelines, and exploratory data analysis.

Create a Sample DataFrame

Let's create a sample DataFrame.


from pyspark.sql import SparkSession
spark = SparkSession.builder.appName("app").getOrCreate()

data = [
    ("James","Sales",3000),
    ("Michael","Sales",4600),
    ("Robert","Sales",4100),
    ("Maria","Finance",3000),
    ("Raman","Finance",3000),
    ("Scott","Finance",3300),
    ("Jen","Finance",3900),
    ("Jeff","Marketing",3000),
    ("Kumar","Marketing",2000)
]

cols = ["Name","Department","Salary"]
df = spark.createDataFrame(data, cols)
df.show()

Extracting First N Rows

When working with large datasets, it’s often useful to take a quick look at the first few rows to understand the structure and content of the data. PySpark offers several ways to extract the first N rows, which we’ll explore below.

Using head()

The head() method retrieves the first N rows of a DataFrame. It returns a list of Row objects, allowing you to access the data as a regular Python list.


# extract top 2 rows
a = df.head(2)
print(a)

# extract top 1 row
a = df.head(1)
print(a)

Yields below output.

Using first()

The first() method retrieves only the very first row from the DataFrame. It’s similar to head(1) but returns a single Row object instead of a list.


# Return first row of DataFrame
a = df.first()
print(a)

# Output:
# Row(Name='James', Department='Sales', Salary=3000)

Using limit(n).collect()

The limit(n) function restricts the number of rows, and collect() converts them into a list of Row objects.


# Use limit() and collect() to return limited number of rows 
a = df.limit(3).collect()
print(a)

# Output:
# [Row(Name='James', Department='Sales', Salary=3000), Row(Name='Michael', Department='Sales', Salary=4600), # Row(Name='Robert', Department='Sales', Salary=4100)]

Extract Top N Rows Per Group

In real-world datasets, we often need to extract the top N rows for each group based on a specific column.
We can achieve this using Window.partitionBy() along with the row_number() function.

For example, let’s get the top 2 employees per Department by Salary.


# Extracting Top N Rows Per Group
from pyspark.sql.window import Window
from pyspark.sql.functions import row_number, col

# Define window partitioned by Department and ordered by Salary (descending)
windowSpec = Window.partitionBy("Department").orderBy(col("Salary").desc())

# Add row_number column
df_topN = df.withColumn("row_num", row_number().over(windowSpec)) \
            .filter(col("row_num") <= 2) \
            .drop("row_num")

df_topN.show()

# Output:
# +-------+----------+------+
# |   Name|Department|Salary|
# +-------+----------+------+
# |    Jen|   Finance|  3900|
# |  Scott|   Finance|  3300|
# |   Jeff| Marketing|  3000|
# |  Kumar| Marketing|  2000|
# |Michael|     Sales|  4600|
# | Robert|     Sales|  4100|
# +-------+----------+------+

Extract Last N Rows

When working with large datasets, it is often useful to quickly inspect the last few rows to gain insights into the most recent entries or the tail end of the data. We can extract the last N rows using the following methods.

Using tail()

The tail() method fetches the last N rows from the DataFrame, returning a list of Row objects similar to head(), but from the end.


# Use tail() to get last n rows
df.tail(3)

# Output:
[Row(Name='Jen', Department='Finance', Salary=3900),
 Row(Name='Jeff', Department='Marketing', Salary=3000),
 Row(Name='Kumar', Department='Marketing', Salary=2000)]

Using collect()

You can also use the collect() method to retrieve all rows from the DataFrame as a list of Rows, and then slice it to get the last N rows.


# Retrieve all columns using collect()
df.collect()[-3:]

# Output:
[Row(Name='Jen', Department='Finance', Salary=3900),
 Row(Name='Jeff', Department='Marketing', Salary=3000),
 Row(Name='Kumar', Department='Marketing', Salary=2000)]

Extracting Last N Rows Per Group

Similarly, to get the last N rows per group, we just need to order the window in ascending order instead of descending.


# Define window partitioned by Department and ordered by Salary (ascending)
windowSpecAsc = Window.partitionBy("Department").orderBy(col("Salary").asc())

# Add row_number column
df_lastN = df.withColumn("row_num", row_number().over(windowSpecAsc)) \
             .filter(col("row_num") <= 2) \
             .drop("row_num")

df_lastN.show()

# Output:
# +------+----------+------+
# |  Name|Department|Salary|
# +------+----------+------+
# | Maria|   Finance|  3000|
# | Raman|   Finance|  3000|
# | Kumar| Marketing|  2000|
# |  Jeff| Marketing|  3000|
# | James|     Sales|  3000|
# |Robert|     Sales|  4100|
# +------+----------+------+

Frequently Asked Questions (FAQs)

first() returns a single Row object.

Use tail(n) instead of collect().

Use rank() or dense_rank() instead of row_number().

Conclusion

In this article, I have explained different ways to extract the first and last N rows from a PySpark DataFrame, both from the entire dataset and within groups.

Simple inspection: head(), first(), limit(), tail()
Grouped extraction: Window.partitionBy().orderBy() with row_number()
Flexible ordering: switch between ascending and descending

By combining these approaches, you can easily retrieve the top or bottom N rows for any analytical or reporting requirement in PySpark.

Happy Learning!!