From 44da2492406b41d1be8f5f19ddf7a869cce1e7b5 Mon Sep 17 00:00:00 2001 From: Arthur Zwiegincew Date: Mon, 10 Jun 2024 11:32:24 -0700 Subject: [PATCH] PR feedback --- src/snowflake/snowpark/_internal/server_connection.py | 2 +- src/snowflake/snowpark/dataframe.py | 4 ++-- src/snowflake/snowpark/mock/_connection.py | 3 ++- src/snowflake/snowpark/session.py | 2 +- 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/snowflake/snowpark/_internal/server_connection.py b/src/snowflake/snowpark/_internal/server_connection.py index 9c97bb95fe1..dd0b45adb38 100644 --- a/src/snowflake/snowpark/_internal/server_connection.py +++ b/src/snowflake/snowpark/_internal/server_connection.py @@ -719,7 +719,7 @@ def ast_query(self, request_id__ast) -> Any: f"/queries/v1/query-request?requestId={request_id}", req ) - def phase1_enabled(self) -> bool: + def is_phase1_enabled(self) -> bool: return os.getenv("SNOWPARK_PHASE_1", False) diff --git a/src/snowflake/snowpark/dataframe.py b/src/snowflake/snowpark/dataframe.py index c1f284c19b8..85842bd9113 100644 --- a/src/snowflake/snowpark/dataframe.py +++ b/src/snowflake/snowpark/dataframe.py @@ -3328,12 +3328,12 @@ def _show_string(self, n: int = 10, max_width: int = 50, **kwargs) -> str: repr.expr.sp_dataframe_show.id.bitfield1 = self._ast_id self._session._ast_batch.eval(repr) - if self._session._conn.phase1_enabled(): + if self._session._conn.is_phase1_enabled(): ast = self._session._ast_batch.flush() res = self._session._conn.ast_query(ast) print(f"AST response: {res}") else: - (_, kwargs["_dataframe_ast"]) = self._session._ast_batch.flush() + _, kwargs["_dataframe_ast"] = self._session._ast_batch.flush() if is_sql_select_statement(query): result, meta = self._session._conn.get_result_and_metadata( diff --git a/src/snowflake/snowpark/mock/_connection.py b/src/snowflake/snowpark/mock/_connection.py index dc066da59cf..7557adcf934 100644 --- a/src/snowflake/snowpark/mock/_connection.py +++ b/src/snowflake/snowpark/mock/_connection.py @@ -800,7 +800,8 @@ def create_coprocessor(self): # It's not necessary to mock this call. pass - def phase1_enabled(self): + def is_phase1_enabled(self): + # We don't yet mock Phase 1. return False diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index c2753a8e5f2..a6b24a0d7b6 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -505,7 +505,7 @@ def __init__( self._runtime_version_from_requirement: str = None # Initialize the server-side session. - if self._conn.phase1_enabled(): + if self._conn.is_phase1_enabled(): self._conn.create_coprocessor() self._ast_batch = AstBatch(self)