PySpark DataFrame Cache and Persist

PySpark Cache and Persist are optimization techniques in DataFrame / Dataset for iterative and interactive PySpark applications to improve the performance of jobs.

Although PySpark provides computation 100 times faster than traditional MapReduce jobs, if you have not designed the jobs to reuse repeating computations, you will see a degradation in performance when dealing with billions or trillions of data. Hence, we may need to examine the stages and utilize optimization techniques as one way to enhance performance.

Using persist() methods, PySpark provides an optimization mechanism to store the intermediate computation of a PySpark DataFrame so they can be reused in subsequent actions.

When you persist a dataset, each node stores its partitioned data in memory and reuses it in other actions on that dataset. And PySpark’s persisted data on nodes is fault-tolerant, meaning if any partition of a Dataset is lost, it will automatically be recomputed using the original transformations that created it.

Advantages of Caching and Persistence of DataFrame

Below are the advantages of using PySpark Cache and Persist methods.

• Cost-efficient – Spark computations are very expensive, hence reusing the computations is used to save cost.
• Time-efficient – Reusing repeated computations saves lots of time.
• Execution time – Saves execution time of the job, and we can perform more jobs on the same cluster.

DataFrame Cache Syntax and Example

PySpark DataFrame or Dataset cache() method by default saves it to the storage level MEMORY_AND_DISK because recomputing the in-memory columnar representation of the underlying table is expensive. Note that this is different from the default cache level of RDD.cache() which is MEMORY_ONLY.

Syntax

# Syntax
DataFrame.cache()

PySpark cache() method in the Dataset class internally calls persist() method that, in turn, uses sparkSession.sharedState.cacheManager.cacheQuery to cache the result set of the DataFrame or Dataset. Let’s look at an example.

Example

First, let’s run some transformations without cache and understand what the performance issue is.

from pyspark.sql import SparkSession
from pyspark.sql.functions import col

spark = SparkSession.builder.appName("example").getOrCreate()

# Read csv with options
df = spark.read.options(header='True', inferSchema='True', delimiter=',') \
    .csv("/apps/PythonTest/sample-zipcodes.csv")

df2 = df.where(col("State") == "PR")
count = df2.count()

df3 = df2.where(col("Zipcode") == 704)
count = df3.count()

DataFrame Persist Syntax and Example

PySpark persist() method is used to store the DataFrame or Dataset to one of the storage levels: MEMORY_ONLY, MEMORY_AND_DISK, MEMORY_ONLY_SER, MEMORY_AND_DISK_SER, DISK_ONLY, MEMORY_ONLY_2, MEMORY_AND_DISK_2 and more.

Caching or persisting of Spark DataFrame or Dataset is a lazy operation, meaning a DataFrame will not be cached until you trigger an action.

Syntax

# persist() Syntax
DataFrame.persist(storageLevel: pyspark.StorageLevel = StorageLevel(True, True, False, True, 1))

PySpark persist has two signatures. The first signature doesn’t take any argument, which by default saves it to MEMORY_AND_DISK storage level, and the second signature, which takes StorageLevel as an argument to store it in different storage levels.

Example

# Persist the DataFrame
dfPersist = df.persist()
dfPersist.show(truncate=False)

Using the second signature, you can save a DataFrame/Dataset to any storage level.

from pyspark import StorageLevel

# Persist the DataFrame
dfPersist = df.persist(StorageLevel.MEMORY_ONLY)
dfPersist.show(truncate=False)

This stores DataFrame/Dataset into memory.

Note that Dataset cache() is an alias for persist(StorageLevel.MEMORY_AND_DISK).

Unpersist Syntax and Example

PySpark automatically monitors every persist() and cache() call you make, and it checks usage on each node and drops persisted data if not used or by using the least-recently-used (LRU) algorithm. You can also manually remove using the unpersist() method. unpersist() marks the Dataset as non-persistent, and removes all blocks for it from memory and disk.

Syntax

# unpersist() Syntax
DataFrame.unpersist(blocking: bool = False) → pyspark.sql.dataframe.DataFrame

Example

# Unpersist the DataFrame
dfPersist = dfPersist.unpersist()

unpersist(Boolean) with a boolean as an argument blocks until all blocks are deleted.

PySpark Persist storage levels

All different storage levels Spark supports are available at pyspark.StorageLevel class. The storage level specifies how and where to persist or cache a Spark DataFrame and Dataset.

• MEMORY_ONLY – This is the default behavior of the RDD cache() method and stores the RDD or DataFrame as deserialized objects to JVM memory. When there is not enough memory available, it will not save all partitions, and these will be re-computed as required. This takes more memory, but unlike RDD, this would be slower than the MEMORY_AND_DISK level as it recomputes the unsaved partitions, and recomputing the in-memory columnar representation of the underlying table is expensive.
• MEMORY_ONLY_SER – This is the same as MEMORY_ONLY, but the difference is it stores RDD as serialized objects to JVM memory. It takes lesser memory (space-efficient) than MEMORY_ONLY as it saves objects as serialized and takes additional CPU cycles to deserialize.
• MEMORY_ONLY_2 – Same as MEMORY_ONLY storage level but replicates each partition to two cluster nodes.
• MEMORY_ONLY_SER_2 – Same as MEMORY_ONLY_SER storage level but replicates each partition to two cluster nodes.
• MEMORY_AND_DISK – This is the default behavior of the DataFrame or Dataset. In this storage level, The DataFrame will be stored in JVM memory as a deserialized object. When required storage is greater than available memory, it stores some of the excess partitions on disk and reads the data from the disk when required. It is slower as there is I/O involved.
• MEMORY_AND_DISK_SER – This is the same as MEMORY_AND_DISK storage level, but it serializes the DataFrame objects in memory and on disk when space is not available.
• MEMORY_AND_DISK_2 – Same as MEMORY_AND_DISK storage level but replicates each partition to two cluster nodes.
• MEMORY_AND_DISK_SER_2 – Same as MEMORY_AND_DISK_SER storage level but replicates each partition to two cluster nodes.
• DISK_ONLY – In this storage level, DataFrame is stored only on disk, and the CPU computation time is high as I/O is involved.
• DISK_ONLY_2 – Same as DISK_ONLY storage level but replicates each partition to two cluster nodes.