SNOW-1782778: Fix Function available memory exhausted

CHANGELOG.md

@@ -10,6 +10,9 @@
 - Added support for `Series.str.center`.
 - Added support for `Series.str.pad`.
 
+#### Bug Fixes
+
+- Fixed a bug that `DataFrame.show` incorrectly fetch all data of the dataframe.
 
 ## 1.26.0 (2024-12-05)
 

src/snowflake/snowpark/_internal/server_connection.py

@@ -491,6 +491,7 @@ def run_query(
         async_post_actions: Optional[List[Query]] = None,
         **kwargs,
     ) -> Union[Dict[str, Any], AsyncJob]:
+        show_limit = kwargs.pop("show_limit", None)
         try:
             # Set SNOWPARK_SKIP_TXN_COMMIT_IN_DDL to True to avoid DDL commands to commit the open transaction
             if is_ddl_on_temp_object:
@@ -524,7 +525,10 @@ def run_query(
             if ignore_results:
                 return {"data": None, "sfqid": results_cursor.sfqid}
             return self._to_data_or_iter(
-                results_cursor=results_cursor, to_pandas=to_pandas, to_iter=to_iter
+                results_cursor=results_cursor,
+                to_pandas=to_pandas,
+                to_iter=to_iter,
+                show_limit=show_limit,
             )
         else:
             return AsyncJob(
@@ -544,7 +548,10 @@ def _to_data_or_iter(
         results_cursor: SnowflakeCursor,
         to_pandas: bool = False,
         to_iter: bool = False,
+        show_limit: int = None,
     ) -> Dict[str, Any]:
+        if show_limit is not None:
+            to_iter = True
         qid = results_cursor.sfqid
         if to_iter:
             new_cursor = results_cursor.connection.cursor()
@@ -580,6 +587,15 @@ def _to_data_or_iter(
             data_or_iter = (
                 iter(results_cursor) if to_iter else results_cursor.fetchall()
             )
+        # return limit number of rows
+        if show_limit is not None:
+            data_or_iter_result = []
+            for i, d in enumerate(data_or_iter):
+                if i >= show_limit:
+                    break
+                else:
+                    data_or_iter_result.append(d)
+            return {"data": data_or_iter_result, "sfqid": qid}
 
         return {"data": data_or_iter, "sfqid": qid}
 

src/snowflake/snowpark/dataframe.py

@@ -4367,6 +4367,7 @@ def _show_string(
                 self.limit(n, _emit_ast=False)._plan, **kwargs
             )
         else:
+            kwargs["show_limit"] = n
             res, meta = self._session._conn.get_result_and_metadata(
                 self._plan, **kwargs
             )