Spark函数内缓存DataFrame的设计合理性与内存管理问题咨询

I have this toy example function returning a cached spark DataFrame (DF):

def foo(df):
    try:
        base = complicated_query(df)
        base.cache()                                    # lazy cache
        base.count()                                    # trigger cache - wrong design???
        num1 = base.withColumn('number', f.lit('1'))
        num2 = base.withColumn('number', f.lit('2'))
        return num1.union(num2)
    finally:
        None
        # base.unpersist()

The purpose of foo is to simply encapsulate temporary variables (DataFrames) which I don't want to have in the outer scope. base is some complicated DF used twice, hence I cache it. My questions are around the explicit count call (suggested by ChatGPT). This is to trigger caching but I feel like it is a wrong design.

• Why do we need count (an action) at this point? What is the gain? The actual caching would happen anyway on the first action to the return value of foo if I didn't call count.
• I noticed that calling foo twice with the same input df has different execution times: first is slow, which is expected since a count is called. But the second is almost instant. Why is that? Surely base is already cached and then count is trivial, but the base reference in the second run is a different reference than that in the first run. How does spark know that it can reuse a cached DF from the first run? (whose memory btw I leaked I guess since after foo exits I cannot unpersist it.)
• Do we really have memory leak when foo exits? How can I unpersist base? Do I have to unpersist at all? I know there is spark.catalog.clearCache() to wipe out all cached DFs but I would like to do this explicitly for base. That's why there is/was a finally clause in the function, to prevent a leak, but that was a failed attempt as in that case I was freeing up the cache before I could even use it...

Can you please help resolve these?

备注：内容来源于stack exchange，提问作者Andras Vanyolos