Repository: databricks/koalas

Issue:
DataFrame.cache()
This would be a super useful function to have, even though it's Spark specific. In terms of documentation, we should create a Spark specific section in frame.rst.




Discussion:
Unlike `.select` which exists in pandas, this function is very specific to spark. How about putting all the spark specific functions under a `.spark` accessor?

It looks like `dask` is not doing that and simply mixes pandas and dask-specific functions, but I am not sure this helps users:
http://docs.dask.org/en/latest/dataframe-api.html#store-dataframes
@rxin, @thunterdb I am taking this up.
hey @shril, have you made some progresses on this?
+1, this is very useful in practice.

Regarding the API, I am thinking that we can add this as a function in the `databricks.koalas` namespace instead of as a method to `DataFrame`. This way, users can write code that works for both pandas and spark dataframes, which helps with writing tests and with transitioning smoothly between koalas and pandas. Then the following code would work with both pandas or koalas:

```py
df = ... # pandas or koalas dataframe
df = ks.cache(df)
# Alternative when chaining:
df = df.pipe(ks.cache)
```

Cons:
 - not as discoverable by the user (will not find it with tab autocomplete)

Also, because caching can consumed resources and should be usually limited in scope, we could implement python guards:

```py
with ks.cache(df):
    mi = df.min()
    ma = df.max()
# Will automatically uncache the dataframe at the end of the block.
```
See an explanation on how to implement here:
http://effbot.org/zone/python-with-statement.htm

Why not both? Discoverability is essential. The other thing is I don't think "you should be able to run exactly the same code in pandas by changing a line of import" is a valid goal. As soon as you merge anything that does approximate count, you will break that goal. So there is really no point in keeping that goal.

@HyukjinKwon I have made some progress. I will update by tonight.
@thunterdb thanks for the input.

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -1509,6 +1509,45 @@ def to_koalas(self):
         else:
             return DataFrame(self)
 
+    def cache(self):
+        """
+        Yields and caches the current DataFrame.
+
+        The Koalas DataFrame is yielded as a protected resource and it's corresponding
+        Spark DataFrame is cached which gets uncached after execution goes of the context.
+
+        Examples
+        --------
+        >>> df = ks.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)],
+        ...                   columns=['dogs', 'cats'])
+        >>> df
+           dogs  cats
+        0   0.2   0.3
+        1   0.0   0.6
+        2   0.6   0.0
+        3   0.2   0.1
+
+        >>> with df.cache() as cached_df:
+        ...     print(cached_df.count())
+        ...
+        dogs    4
+        cats    4
+        dtype: int64
+
+        >>> df = df.cache()
+        >>> df.to_pandas().mean(axis=1)
+        0    0.25
+        1    0.30
+        2    0.30
+        3    0.15
+        dtype: float64
+
+        To uncache the dataframe, use `unpersist` function
+
+        >>> df.unpersist()
+        """
+        return _CachedDataFrame(self._sdf, self._metadata)
+
     def to_table(self, name: str, format: Optional[str] = None, mode: str = 'error',
                  partition_cols: Union[str, List[str], None] = None,
                  **options):
@@ -3603,3 +3642,37 @@ def _reduce_spark_multi(sdf, aggs):
     l2 = list(row)
     assert len(l2) == len(aggs), (row, l2)
     return l2
+
+
+class _CachedDataFrame(DataFrame):
+    """
+    Cached Koalas DataFrame, which corresponds to Pandas DataFrame logically, but internally
+    it caches the corresponding Spark DataFrame.
+    """
+    def __init__(self, sdf, metadata):
+        self._cached = sdf.cache()
+        super(_CachedDataFrame, self).__init__(self._cached, index=metadata)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self.unpersist()
+
+    def unpersist(self):
+        """
+        The `unpersist` function is used to uncache the Koalas DataFrame when it
+        is not used with `with` statement.
+
+        Examples
+        --------
+        >>> df = ks.DataFrame([(.2, .3), (.0, .6), (.6, .0), (.2, .1)],
+        ...                   columns=['dogs', 'cats'])
+        >>> df = df.cache()
+
+        To uncache the dataframe, use `unpersist` function
+
+        >>> df.unpersist()
+        """
+        if self._cached.is_cached:
+            self._cached.unpersist()


Test patch: