fixing dataframe_extensions

CHANGE_LOG.txt

@@ -207,4 +207,8 @@ Version 0.0.39
 Version 0.0.40
 --------------
 - unpinning snowpark dependency
-- removing deprecated functionality already provided in snowpark
+- removing deprecated functionality already provided in snowpark
+
+Version 0.0.41
+--------------
+- Fixing issue #40 when using dataframe_extensions to display a DataFrame inside a notebook

snowpark_extensions/__init__.py

@@ -8,7 +8,7 @@
 from .session_builder_extensions import *
 from .types_extensions import *
 from .logging_utils import logged, Tag
-from .wheel_loader import add_wheels, load
+
 
 
 def register_sql_magic():

snowpark_extensions/dataframe_extensions.py

@@ -47,18 +47,14 @@ def _repr_html_(self):
         else:
             from IPython.display import display
         try:
-            count = self._cached_rowcount if hasattr(self,"_cached_rowcount") else self.count()
-            self.count()
-            if count == 0:
+            data_to_display = self.limit(rows_limit).collect()
+            if len(data_to_display) == 0:
                 return "No rows to display"
-            elif count == 1:
-                df = pd.DataFrame.from_records([x.as_dict() for x in self.collect()])
-            elif count > rows_limit:
-                print(f"There are {count} rows. Showing only {rows_limit}. Change DataFrame.__rows_limit value to display more rows")
-                df = self.limit(rows_limit).to_pandas()
             else:
-                df = self.to_pandas()
-            display(df)
+                df = pd.DataFrame.from_records([x.as_dict() for x in data_to_display])
+                if len(data_to_display) >= rows_limit:
+                    print(f"Showing only {rows_limit}. Change DataFrame.__rows_limit value to display more rows")
+                display(df)
             return ""
         except Exception as ex:
                 return str(ex)