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).
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.
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:
- What is the most efficient way of caching a sample of data from a large dataframe?
- 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?