Repository: databricks/koalas

Issue:
repr and repr_html should not say "1000 rows" when there are more than 1000 rows
See repro:

```python
import databricks.koalas as ks
ks.range(2000)

Out[3]: 
      id
0      0
1      1
2      2
3      3
4      4
5      5
6      6
7      7
8      8
9      9
10    10
11    11
12    12
13    13
14    14
15    15
16    16
17    17
18    18
19    19
20    20
21    21
22    22
23    23
24    24
25    25
26    26
27    27
28    28
29    29
..   ...
970  970
971  971
972  972
973  973
974  974
975  975
976  976
977  977
978  978
979  979
980  980
981  981
982  982
983  983
984  984
985  985
986  986
987  987
988  988
989  989
990  990
991  991
992  992
993  993
994  994
995  995
996  996
997  997
998  998
999  999
[1000 rows x 1 columns]
```

The same problem exists for _repr_html_ as well. The reason this happens is we take the first 1000 rows and convert it to a pandas DataFrame, and leverage pandas' repr and repr_html methods to show them. We need to intercept the output and replace the "1000 rows" with something like "showing only the first 1000 rows".



Discussion:

Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -17,6 +17,7 @@
 """
 A wrapper class for Spark DataFrame to behave similar to pandas DataFrame.
 """
+import re
 import warnings
 from functools import partial, reduce
 from typing import Any, List, Tuple, Union
@@ -40,6 +41,14 @@
 from databricks.koalas.typedef import infer_pd_series_spark_type
 
 
+# These regular expression patterns are complied and defined here to avoid to compile the same
+# pattern every time it is used in _repr_ and _repr_html_ in DataFrame.
+# Two patterns basically seek the footer string from Pandas'
+REPR_PATTERN = re.compile(r"\n\n\[(?P<rows>[0-9]+) rows x (?P<columns>[0-9]+) columns\]$")
+REPR_HTML_PATTERN = re.compile(
+    r"\n\<p\>(?P<rows>[0-9]+) rows × (?P<columns>[0-9]+) columns\<\/p\>\n\<\/div\>$")
+
+
 class DataFrame(_Frame):
     """
     Koala DataFrame that corresponds to Pandas DataFrame logically. This holds Spark DataFrame
@@ -2355,10 +2364,35 @@ def _pd_getitem(self, key):
         raise NotImplementedError(key)
 
     def __repr__(self):
-        return repr(self.head(max_display_count).to_pandas())
+        pdf = self.head(max_display_count + 1).to_pandas()
+        pdf_length = len(pdf)
+        repr_string = repr(pdf.iloc[:max_display_count])
+        if pdf_length > max_display_count:
+            match = REPR_PATTERN.search(repr_string)
+            if match is not None:
+                nrows = match.group("rows")
+                ncols = match.group("columns")
+                footer = ("\n\n[Showing only the first {nrows} rows x {ncols} columns]"
+                          .format(nrows=nrows, ncols=ncols))
+                return REPR_PATTERN.sub(footer, repr_string)
+        return repr_string
 
     def _repr_html_(self):
-        return self.head(max_display_count).to_pandas()._repr_html_()
+        pdf = self.head(max_display_count + 1).to_pandas()
+        pdf_length = len(pdf)
+        repr_html = pdf[:max_display_count]._repr_html_()
+        if pdf_length > max_display_count:
+            match = REPR_HTML_PATTERN.search(repr_html)
+            if match is not None:
+                nrows = match.group("rows")
+                ncols = match.group("columns")
+                by = chr(215)
+                footer = ('\n<p>Showing only the first {rows} rows {by} {cols} columns</p>\n</div>'
+                          .format(rows=nrows,
+                                  by=by,
+                                  cols=ncols))
+                return REPR_HTML_PATTERN.sub(footer, repr_html)
+        return repr_html
 
     def __getitem__(self, key):
         return self._pd_getitem(key)
diff --git a/databricks/koalas/series.py b/databricks/koalas/series.py
--- a/databricks/koalas/series.py
+++ b/databricks/koalas/series.py
@@ -17,6 +17,7 @@
 """
 A wrapper class for Spark Column to behave similar to pandas Series.
 """
+import re
 import inspect
 from functools import partial, wraps
 from typing import Any
@@ -39,6 +40,12 @@
 from databricks.koalas.utils import validate_arguments_and_invoke_function
 
 
+# This regular expression pattern is complied and defined here to avoid to compile the same
+# pattern every time it is used in _repr_ in Series.
+# This pattern basically seeks the footer string from Pandas'
+REPR_PATTERN = re.compile(r"Length: (?P<length>[0-9]+)")
+
+
 def _column_op(f):
     """
     A decorator that wraps APIs taking/returning Spark Column so that Koalas Series can be
@@ -1086,7 +1093,18 @@ def __str__(self):
         return self._pandas_orig_repr()
 
     def __repr__(self):
-        return repr(self.head(max_display_count).to_pandas())
+        pser = self.head(max_display_count + 1).to_pandas()
+        pser_length = len(pser)
+        repr_string = repr(pser.iloc[:max_display_count])
+        if pser_length > max_display_count:
+            rest, prev_footer = repr_string.rsplit("\n", 1)
+            match = REPR_PATTERN.search(prev_footer)
+            if match is not None:
+                length = match.group("length")
+                footer = ("\n{prev_footer}\nShowing only the first {length}"
+                          .format(length=length, prev_footer=prev_footer))
+                return rest + footer
+        return repr_string
 
     def __dir__(self):
         if not isinstance(self.schema, StructType):


Test patch:
diff --git a/databricks/koalas/tests/test_dataframe.py b/databricks/koalas/tests/test_dataframe.py
--- a/databricks/koalas/tests/test_dataframe.py
+++ b/databricks/koalas/tests/test_dataframe.py
@@ -69,11 +69,6 @@ def test_dataframe(self):
 
         self.assertEqual(ddf.a.notnull().alias("x").name, "x")
 
-    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)
@@ -81,11 +76,6 @@ def test_repr_cache_invalidation(self):
         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)
diff --git a/databricks/koalas/tests/test_repr.py b/databricks/koalas/tests/test_repr.py
new file mode 100644
--- /dev/null
+++ b/databricks/koalas/tests/test_repr.py
@@ -0,0 +1,45 @@
+#
+# Copyright (C) 2019 Databricks, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from databricks import koalas as ks
+from databricks.koalas.generic import max_display_count
+from databricks.koalas.testing.utils import ReusedSQLTestCase
+
+
+class ReprTests(ReusedSQLTestCase):
+
+    def test_repr_dataframe(self):
+        kdf = ks.range(max_display_count)
+        self.assertTrue("Showing only the first" not in repr(kdf))
+        self.assert_eq(repr(kdf), repr(kdf.to_pandas()))
+
+        kdf = ks.range(max_display_count + 1)
+        self.assertTrue("Showing only the first" in repr(kdf))
+
+    def test_repr_series(self):
+        kser = ks.range(max_display_count).id
+        self.assertTrue("Showing only the first" not in repr(kser))
+        self.assert_eq(repr(kser), repr(kser.to_pandas()))
+
+        kser = ks.range(max_display_count + 1).id
+        self.assertTrue("Showing only the first" in repr(kser))
+
+    def test_html_repr(self):
+        kdf = ks.range(max_display_count)
+        self.assertTrue("Showing only the first" not in kdf._repr_html_())
+        self.assertEqual(kdf._repr_html_(), kdf.to_pandas()._repr_html_())
+
+        kdf = ks.range(max_display_count + 1)
+        self.assertTrue("Showing only the first" in kdf._repr_html_())
diff --git a/databricks/koalas/tests/test_series.py b/databricks/koalas/tests/test_series.py
--- a/databricks/koalas/tests/test_series.py
+++ b/databricks/koalas/tests/test_series.py
@@ -44,11 +44,6 @@ 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']