From 18eafb3497a72047f65cdc31845b0d90a03660e6 Mon Sep 17 00:00:00 2001
From: Nick Crews <nicholas.b.crews@gmail.com>
Date: Thu, 21 Nov 2024 07:34:50 -0800
Subject: [PATCH] feat(repr): show col and row count in rich repr

---
 ibis/config.py               |  4 ++++
 ibis/expr/types/generic.py   |  6 +++++-
 ibis/expr/types/pretty.py    | 13 ++++++++++++-
 ibis/expr/types/relations.py |  4 ++++
 4 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/ibis/config.py b/ibis/config.py
index 151cabbb2d85..6c620a077d46 100644
--- a/ibis/config.py
+++ b/ibis/config.py
@@ -67,6 +67,9 @@ class Interactive(Config):
         Maximum depth for nested data types.
     show_types : bool
         Show the inferred type of value expressions in the interactive repr.
+    show_count: bool
+        For Columns and Tables, show the row count. This can be computationally
+        expensive and slow.
 
     """
 
@@ -76,6 +79,7 @@ class Interactive(Config):
     max_string: int = 80
     max_depth: int = 1
     show_types: bool = True
+    show_count: bool = True
 
 
 class Repr(Config):
diff --git a/ibis/expr/types/generic.py b/ibis/expr/types/generic.py
index 17559746d7aa..9a0a54dcae4d 100644
--- a/ibis/expr/types/generic.py
+++ b/ibis/expr/types/generic.py
@@ -1366,6 +1366,7 @@ def preview(
         max_length: int | None = None,
         max_string: int | None = None,
         max_depth: int | None = None,
+        show_count: bool | None = None,
         console_width: int | float | None = None,
     ) -> rich.table.Table:
         """Print a subset as a single-column Rich Table.
@@ -1385,6 +1386,8 @@ def preview(
             Maximum length for pretty-printed strings.
         max_depth
             Maximum depth for nested data types.
+        show_count
+            Show the row count. This can be computationally expensive and slow.
         console_width
             Width of the console in characters. If not specified, the width
             will be inferred from the console.
@@ -1412,10 +1415,11 @@ def preview(
             max_length=max_length,
             max_string=max_string,
             max_depth=max_depth,
+            show_count=show_count,
         )
         overrides = {k: v for k, v in overrides.items() if v is not None}
         options = ibis.options.repr.interactive.copy(**overrides)
-        return to_rich(self, options, console_width=console_width)
+        return to_rich(self, options=options, console_width=console_width)
 
     def __pyarrow_result__(
         self,
diff --git a/ibis/expr/types/pretty.py b/ibis/expr/types/pretty.py
index fea1d3a24b62..46d838cc1555 100644
--- a/ibis/expr/types/pretty.py
+++ b/ibis/expr/types/pretty.py
@@ -292,6 +292,8 @@ def _to_rich_table(
     options: Interactive,
     console_width: int | float | None = None,
 ) -> rich.table.Table:
+    from ibis.expr import types as ir
+
     console_width = console_width or float("inf")
     max_rows = options.max_rows
     max_columns = options.max_columns
@@ -411,7 +413,16 @@ def _to_rich_table(
             if not next_flex_cols:
                 break
 
-    rich_table = rich.table.Table(padding=(0, 1, 0, 1))
+    if options.show_count:
+        # use underscore to be friendly to i18n and python REPL
+        nrows = f"{table.count().execute():_}"
+    else:
+        nrows = "…"
+    if isinstance(tablish, ir.Table):
+        dims = f"{orig_ncols:_} cols by {nrows} rows"
+    else:
+        dims = None if nrows == "…" else f"{nrows} rows"
+    rich_table = rich.table.Table(title=dims, padding=(0, 1, 0, 1))
 
     # Configure the columns on the rich table.
     for name, dtype, _, max_width in col_info:
diff --git a/ibis/expr/types/relations.py b/ibis/expr/types/relations.py
index 6efab7d20580..1992889ece19 100644
--- a/ibis/expr/types/relations.py
+++ b/ibis/expr/types/relations.py
@@ -507,6 +507,7 @@ def preview(
         max_length: int | None = None,
         max_string: int | None = None,
         max_depth: int | None = None,
+        show_count: bool | None = None,
         console_width: int | float | None = None,
     ) -> RichTable:
         """Return a subset as a Rich Table.
@@ -528,6 +529,8 @@ def preview(
             Maximum length for pretty-printed strings
         max_depth
             Maximum depth for nested data types
+        show_count
+            Show the row count. This can be computationally expensive and slow.
         console_width
             Width of the console in characters. If not specified, the width
             will be inferred from the console.
@@ -565,6 +568,7 @@ def preview(
             max_length=max_length,
             max_string=max_string,
             max_depth=max_depth,
+            show_count=show_count,
         )
         overrides = {k: v for k, v in overrides.items() if v is not None}
         options = ibis.options.repr.interactive.copy(**overrides)