PySpark : Understanding the ‘take’ Action in PySpark with Examples. [Retrieves a specified number of elements from the beginning of an RDD or DataFrame]

In this article, we will focus on the ‘take’ action, which is commonly used in PySpark operations. We’ll provide a brief explanation of the ‘take’ action, followed by a simple example to help you understand its usage.

What is the ‘take’ Action in PySpark?

The ‘take’ action in PySpark retrieves a specified number of elements from the beginning of an RDD (Resilient Distributed Dataset) or DataFrame. It is an action operation, which means it triggers the execution of any previous transformations on the data, returning the result to the driver program. This operation is particularly useful for previewing the contents of an RDD or DataFrame without having to collect all the elements, which can be time-consuming and memory-intensive for large datasets.

Syntax:

take(num)

Where num is the number of elements to retrieve from the RDD or DataFrame.

Simple Example

Let’s go through a simple example using the ‘take’ action in PySpark. First, we’ll create a PySpark RDD and then use the ‘take’ action to retrieve a specified number of elements.

RDD Version

Step 1: Start a PySpark session

Before starting with the example, you’ll need to start a PySpark session:

from pyspark.sql import SparkSession
spark = SparkSession.builder \
    .appName("Understanding the 'take' action in PySpark") \
    .getOrCreate()

data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
rdd = spark.sparkContext.parallelize(data)

data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
rdd = spark.sparkContext.parallelize(data)

first_five_elements = rdd.take(5)
print("The first five elements of the RDD are:", first_five_elements)

The first five elements of the RDD are: [1, 2, 3, 4, 5]

from pyspark.sql import SparkSession
spark = SparkSession.builder \
    .appName("Understanding the 'take' action in PySpark with DataFrames") \
    .getOrCreate()
from pyspark.sql import Row
data = [
    Row(name="Alice", age=30, city="New York"),
    Row(name="Bob", age=28, city="San Francisco"),
    Row(name="Cathy", age=25, city="Los Angeles"),
    Row(name="David", age=32, city="Chicago"),
    Row(name="Eva", age=29, city="Seattle")
]
schema = "name STRING, age INT, city STRING"
df = spark.createDataFrame(data, schema=schema)
first_three_rows = df.take(3)
print("The first three rows of the DataFrame are:")
for row in first_three_rows:
    print(row)

The first three rows of the DataFrame are:
Row(name='Alice', age=30, city='New York')
Row(name='Bob', age=28, city='San Francisco')
Row(name='Cathy', age=25, city='Los Angeles')