Skip to content

Commit

Permalink
fix error
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-yzou committed Nov 8, 2024
1 parent 07c955d commit dcfaf21
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 5 deletions.
2 changes: 1 addition & 1 deletion src/snowflake/snowpark/modin/plugin/_internal/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -273,7 +273,7 @@ def _create_read_only_table(
readonly_table_name = (
f"{random_name_for_temp_object(TempObjectType.TABLE)}{READ_ONLY_TABLE_SUFFIX}"
)
use_scoped_temp_table = session._use_scoped_temp_objects
use_scoped_temp_table = session._use_scoped_temp_read_only_table
# If we need to materialize into a temp table our create table expression
# needs to be SELECT * FROM (object).
if materialize_into_temp_table:
Expand Down
9 changes: 9 additions & 0 deletions src/snowflake/snowpark/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,6 +238,10 @@
_PYTHON_SNOWPARK_ENABLE_THREAD_SAFE_SESSION = (
"PYTHON_SNOWPARK_ENABLE_THREAD_SAFE_SESSION"
)
# Flag for controlling the usage of scoped temp read only table.
_PYTHON_SNOWPARK_ENABLE_SCOPED_TEMP_READ_ONLY_TABLE = (
"PYTHON_SNOWPARK_ENABLE_SCOPED_TEMP_READ_ONLY_TABLE"
)
# The complexity score lower bound is set to match COMPILATION_MEMORY_LIMIT
# in Snowflake. This is the limit where we start seeing compilation errors.
DEFAULT_COMPLEXITY_SCORE_LOWER_BOUND = 10_000_000
Expand Down Expand Up @@ -541,6 +545,11 @@ def __init__(
_PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS_STRING, True
)
)
self._use_scoped_temp_read_only_table: bool = (
self._conn._get_client_side_session_parameter(
_PYTHON_SNOWPARK_ENABLE_SCOPED_TEMP_READ_ONLY_TABLE, False
)
)
self._file = FileOperation(self)
self._lineage = Lineage(self)
self._sql_simplifier_enabled: bool = (
Expand Down
8 changes: 4 additions & 4 deletions tests/integ/modin/io/test_read_snowflake.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,10 @@

@pytest.fixture(params=paramList)
def setup_use_scoped_object(request, session):
use_scoped_objects = session._use_scoped_temp_objects
session._use_scoped_temp_objects = request.param
use_scoped_objects = session._use_scoped_temp_read_only_table
session._use_scoped_temp_read_only_table = request.param
yield
session._use_scoped_temp_objects = use_scoped_objects
session._use_scoped_temp_read_only_table = use_scoped_objects


def read_snowflake_and_verify_snapshot_creation(
Expand Down Expand Up @@ -85,7 +85,7 @@ def read_snowflake_and_verify_snapshot_creation(
assert len(query_history.queries) == 1

# test if the scoped snapshot is created
scoped_pattern = " SCOPED " if session._use_scoped_temp_objects else " "
scoped_pattern = " SCOPED " if session._use_scoped_temp_read_only_table else " "
table_create_sql = query_history.queries[-1].sql_text
table_create_pattern = f"CREATE OR REPLACE{scoped_pattern}TEMPORARY READ ONLY TABLE SNOWPARK_TEMP_TABLE_[0-9A-Z]+.*{READ_ONLY_TABLE_SUFFIX}.*"
assert re.match(table_create_pattern, table_create_sql) is not None
Expand Down

0 comments on commit dcfaf21

Please sign in to comment.