diff --git a/src/snowflake/snowpark/__init__.py b/src/snowflake/snowpark/__init__.py index d2b2820575..231aad77e8 100644 --- a/src/snowflake/snowpark/__init__.py +++ b/src/snowflake/snowpark/__init__.py @@ -67,3 +67,36 @@ WhenNotMatchedClause, ) from snowflake.snowpark.window import Window, WindowSpec + + +def register_sql_magic(): + try: + import IPython + from IPython.core.magic import register_cell_magic + + def sql(line, cell): + user_ns = IPython.get_ipython().user_ns + if "session" in user_ns: + session = user_ns["session"] + # the following will allow using variables + # for example session.sql("select {variable1}") + formatted_sql = cell.format(**user_ns) + name = None + line = (line or "").strip() + if line: + name = line.strip().split(" ")[0] + df = session.sql(formatted_sql) + if name: + user_ns[name] = df + else: + user_ns["__df"] = df + return df + else: + return "No session was found. Define a Snowpark Session object with the name 'session'" + + register_cell_magic(sql) + except Exception: + pass + + +register_sql_magic() diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index f2a79a6520..f33a66c333 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -3772,6 +3772,49 @@ def convert(col: ColumnOrName) -> Expression: exprs = [convert(col) for col in parse_positional_args_to_list(*cols)] return exprs + # class property to control the number of rows to diplay + __rows = 50 + + def _repr_html_(self): + """Render a rich DataFrame representation in HTML format. + + This method is used for pretty printing DataFrame in Jupyter + notebook or similar environments that support HTML rendering.""" + import IPython + import pandas as pd + + rows_limit = DataFrame.__rows + if "display" in globals(): + display = globals()["display"] + elif "display" in IPython.get_ipython().user_ns: + display = IPython.get_ipython().user_ns["display"] + else: + from IPython.display import display + try: + count = self.count() + if count == 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) + return "" + except Exception as ex: + return str(ex) + + def _ipython_key_completions_(self) -> List[str]: + """Provide custom key completions for the DataFrame in IPython or Jupyter Notebook. + + This method is used to supply additional keys for tab-completion support in + IPython or Jupyter Notebook when interacting with the DataFrame.""" + return self.columns + where = filter # Add the following lines so API docs have them diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index f36c30e8d3..7beffe7ef0 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -2284,4 +2284,22 @@ def _get_client_side_session_parameter(self, name: str, default_value: Any) -> A else default_value ) + def __repr__(self): + """ + Return a string representation of the Session object.""" + snowflake_environment = self.sql("select current_version()").collect() + return ( + "\n" + + f"role : {self.get_current_role()}\n" + + f"warehouse : {self.get_current_warehouse()}\n" + + f"database : {self.get_current_database()}\n" + + f"schema : {self.get_current_schema()}\n" + + f"db.version : {snowflake_environment[0][0]}\n" + + f"snowpark.version : {get_version()}\n" + + f"os.name : {get_os_name()}\n" + + f"python.connector.version : {get_connector_version()}\n" + + f"python.connector.session.id : {self._session_id}\n" + + f"query_tags : {self.query_tag}" + ) + createDataFrame = create_dataframe