Adding .cache() to (Py)Spark dataframe ignores .limit() and caches complete dataframe

Viewed 196

I'm writing up a pipeline for a very large dataframe. In order to rapidly prototype, I'm trying to cache a small amount of samples from the dataframe.

My expectation was that I could efficiently achieve this by executing the code below.

df.limit(1).cache().collect()

However, as shown in the picture below, this creates two stages of which the first one caches the complete dataframe (toy example of 99.3gb).

Spark stages with cache()

This is surprising, because when I remove the cache method (leaving df.limit(1).collect()), spark spawns a collect stage that is only applied on a single partition in my dataframe. See picture below.

Spark stages without cache()

This same different is also observed when comparing the stages when combining limit with show/collect/take (stage with one partition) vs count (stage with full dataframe).

Although I can speed up this cache by filtering out a single partition manually, I'm still stuck with the following two questions:

  1. What is the most efficient way of caching a sample of data from a large dataframe?
  2. Why is .cache() caching the complete dataframe when a limit is placed on the query, while the same query without cache() doesn't touch the complete dataframe?
0 Answers
Related