repr and repr_html improvements (#258)

Author: rxin

databricks/koalas/frame.py

@@ -733,6 +733,13 @@ def to_koalas(self):
            __index_level_0__  col1  col2
         0                  0     1     3
         1                  1     2     4
+
+        Calling to_koalas on a Koalas DataFrame simply returns itself.
+
+        >>> df.to_koalas()
+           col1  col2
+        0     1     3
+        1     2     4
         """
         if isinstance(self, DataFrame):
             return self
@@ -1298,33 +1305,45 @@ def sort_values(self, by, ascending=True, inplace=False, na_position='last'):
         Examples
         --------
         >>> df = ks.DataFrame({
-        ...     'col1': ['A', 'A', 'B', None, 'D', 'C'],
-        ...     'col2': [2, 1, 9, 8, 7, 4],
-        ...     'col3': [0, 1, 9, 4, 2, 3],
+        ...     'col1': ['A', 'B', None, 'D', 'C'],
+        ...     'col2': [2, 9, 8, 7, 4],
+        ...     'col3': [0, 9, 4, 2, 3],
         ... })
         >>> df
            col1  col2  col3
         0     A     2     0
-        1     A     1     1
-        2     B     9     9
-        3  None     8     4
-        4     D     7     2
-        5     C     4     3
+        1     B     9     9
+        2  None     8     4
+        3     D     7     2
+        4     C     4     3
 
         Sort by col1
 
         >>> df.sort_values(by=['col1'])
            col1  col2  col3
         0     A     2     0
-        1     A     1     1
-        2     B     9     9
-        5     C     4     3
-        4     D     7     2
-        3  None     8     4
+        1     B     9     9
+        4     C     4     3
+        3     D     7     2
+        2  None     8     4
 
+        Sort Descending
+
+        >>> df.sort_values(by='col1', ascending=False)
+           col1  col2  col3
+        3     D     7     2
+        4     C     4     3
+        1     B     9     9
+        0     A     2     0
+        2  None     8     4
 
         Sort by multiple columns
 
+        >>> df = ks.DataFrame({
+        ...     'col1': ['A', 'A', 'B', None, 'D', 'C'],
+        ...     'col2': [2, 1, 9, 8, 7, 4],
+        ...     'col3': [0, 1, 9, 4, 2, 3],
+        ... })
         >>> df.sort_values(by=['col1', 'col2'])
            col1  col2  col3
         1     A     1     1
@@ -1333,17 +1352,6 @@ def sort_values(self, by, ascending=True, inplace=False, na_position='last'):
         5     C     4     3
         4     D     7     2
         3  None     8     4
-
-        Sort Descending
-
-        >>> df.sort_values(by='col1', ascending=False)
-           col1  col2  col3
-        4     D     7     2
-        5     C     4     3
-        2     B     9     9
-        0     A     2     0
-        1     A     1     1
-        3  None     8     4
         """
         if isinstance(by, string_types):
             by = [by]
@@ -1496,7 +1504,10 @@ def _pd_getitem(self, key):
         raise NotImplementedError(key)
 
     def __repr__(self):
-        return repr(self.toPandas())
+        return repr(self.head(max_display_count).to_pandas())
+
+    def _repr_html_(self):
+        return self.head(max_display_count).to_pandas()._repr_html_()
 
     def __getitem__(self, key):
         return self._pd_getitem(key)
@@ -1541,9 +1552,6 @@ def __dir__(self):
         fields = [f for f in self._sdf.schema.fieldNames() if ' ' not in f]
         return super(DataFrame, self).__dir__() + fields
 
-    def _repr_html_(self):
-        return self.head(max_display_count).toPandas()._repr_html_()
-
     @classmethod
     def _validate_axis(cls, axis=0):
         if axis not in (0, 1, 'index', 'columns', None):

databricks/koalas/generic.py

@@ -164,7 +164,7 @@ def groupby(self, by):
                         'got [%s]' % (df_or_s,))
 
     def compute(self):
-        """Alias of `toPandas()` to mimic dask for easily porting tests."""
+        """Alias of `to_pandas()` to mimic dask for easily porting tests."""
         return self.toPandas()
 
     @staticmethod

databricks/koalas/series.py

@@ -731,7 +731,7 @@ def __str__(self):
         return self._pandas_orig_repr()
 
     def __repr__(self):
-        return repr(self.head(max_display_count).toPandas())
+        return repr(self.head(max_display_count).to_pandas())
 
     def __dir__(self):
         if not isinstance(self.schema, StructType):

databricks/koalas/tests/test_dataframe.py

@@ -20,6 +20,7 @@
 import pandas as pd
 
 from databricks import koalas
+from databricks.koalas.generic import max_display_count
 from databricks.koalas.testing.utils import ReusedSQLTestCase, SQLTestUtils
 from databricks.koalas.exceptions import PandasNotImplementedError
 from databricks.koalas.missing.frame import _MissingPandasLikeDataFrame
@@ -68,8 +69,31 @@ def test_dataframe(self):
 
         self.assertEqual(ddf.a.notnull().alias("x").name, "x")
 
-    def test_empty_dataframe(self):
+    def test_repr(self):
+        # Make sure we only fetch max_display_count
+        self.assertEqual(koalas.range(1001).__repr__(),
+                         koalas.range(max_display_count).__repr__())
+
+    def test_repr_cache_invalidation(self):
+        # If there is any cache, inplace operations should invalidate it.
+        df = koalas.range(10)
+        df.__repr__()
+        df['a'] = df['id']
+        self.assertEqual(df.__repr__(), df.to_pandas().__repr__())
+
+    def test_repr_html(self):
+        # Make sure we only fetch max_display_count
+        self.assertEqual(koalas.range(1001)._repr_html_(),
+                         koalas.range(max_display_count)._repr_html_())
+
+    def test_repr_html_cache_invalidation(self):
+        # If there is any cache, inplace operations should invalidate it.
+        df = koalas.range(10)
+        df._repr_html_()
+        df['a'] = df['id']
+        self.assertEqual(df._repr_html_(), df.to_pandas()._repr_html_())
 
+    def test_empty_dataframe(self):
         pdf = pd.DataFrame({'a': pd.Series([], dtype='i1'),
                             'b': pd.Series([], dtype='str')})
 

databricks/koalas/tests/test_series.py

@@ -21,6 +21,7 @@
 
 from databricks import koalas
 from databricks.koalas import Series
+from databricks.koalas.generic import max_display_count
 from databricks.koalas.testing.utils import ReusedSQLTestCase, SQLTestUtils
 from databricks.koalas.exceptions import PandasNotImplementedError
 from databricks.koalas.missing.series import _MissingPandasLikeSeries
@@ -43,6 +44,18 @@ def test_series(self):
 
         # TODO: self.assert_eq(d + 1, pdf + 1)
 
+    def test_repr(self):
+        # Make sure we only fetch max_display_count
+        self.assertEqual(koalas.range(1001)['id'].__repr__(),
+                         koalas.range(max_display_count)['id'].__repr__())
+
+    def test_repr_cache_invalidation(self):
+        # If there is any cache, inplace operations should invalidate it.
+        s = koalas.range(10)['id']
+        s.__repr__()
+        s.rename('a', inplace=True)
+        self.assertEqual(s.__repr__(), s.rename("a").__repr__())
+
     def test_empty_series(self):
         a = pd.Series([], dtype='i1')
         b = pd.Series([], dtype='str')

databricks/koalas/tests/test_utils.py

@@ -17,7 +17,9 @@
 import pandas as pd
 
 from databricks.koalas.testing.utils import ReusedSQLTestCase, SQLTestUtils
-from databricks.koalas.utils import validate_arguments_and_invoke_function
+from databricks.koalas.utils import lazy_property, validate_arguments_and_invoke_function
+
+some_global_variable = 0
 
 
 class UtilsTest(ReusedSQLTestCase, SQLTestUtils):
@@ -43,3 +45,20 @@ def test_validate_arguments_and_invoke_function(self):
         # to a non-default value
         with self.assertRaises(TypeError):
             self.to_html(unsupported_param=1)
+
+    def test_lazy_property(self):
+        obj = TestClassForLazyProp()
+        # If lazy prop is not working, the second test would fail (because it'd be 2)
+        self.assert_eq(obj.lazy_prop, 1)
+        self.assert_eq(obj.lazy_prop, 1)
+
+
+class TestClassForLazyProp:
+
+    def __init__(self):
+        self.some_variable = 0
+
+    @lazy_property
+    def lazy_prop(self):
+        self.some_variable += 1
+        return self.some_variable

databricks/koalas/utils.py

@@ -73,3 +73,19 @@ def validate_arguments_and_invoke_function(pobj: Union[pd.DataFrame, pd.Series],
 
     args['self'] = pobj
     return pandas_func(**args)
+
+
+def lazy_property(fn):
+    """
+    Decorator that makes a property lazy-evaluated.
+
+    Copied from https://stevenloria.com/lazy-properties/
+    """
+    attr_name = '_lazy_' + fn.__name__
+
+    @property
+    def _lazy_property(self):
+        if not hasattr(self, attr_name):
+            setattr(self, attr_name, fn(self))
+        return getattr(self, attr_name)
+    return _lazy_property