SNOW-1546396: use scoped temp tables for large query breakdown (#2197)

<!---
Please answer these questions before creating your pull request. Thanks!
--->

1. Which Jira issue is this PR addressing? Make sure that there is an
accompanying issue to your PR.

   <!---
   In this section, please add a Snowflake Jira issue number.
   
Note that if a corresponding GitHub issue exists, you should still
include
   the Snowflake Jira issue number. For example, for GitHub issue
#1400, you should
   add "SNOW-1335071" here.
    --->

   Fixes SNOW-1546396

2. Fill out the following pre-review checklist:

- [ ] I am adding a new automated test(s) to verify correctness of my
new code
- [ ] If this test skips Local Testing mode, I'm requesting review from
@snowflakedb/local-testing
   - [ ] I am adding new logging messages
   - [ ] I am adding a new telemetry message
   - [ ] I am adding new credentials
   - [ ] I am adding a new dependency
- [ ] If this is a new feature/behavior, I'm adding the Local Testing
parity changes.

3. Please describe how your code solves the related issue.

   Use scoped temp table for large query breakdown.

src/snowflake/snowpark/_internal/analyzer/analyzer_utils.py

@@ -856,6 +856,9 @@ def create_table_as_select_statement(
     max_data_extension_time: Optional[int] = None,
     change_tracking: Optional[bool] = None,
     copy_grants: bool = False,
+    *,
+    use_scoped_temp_objects: bool = False,
+    is_generated: bool = False,
 ) -> str:
     column_definition_sql = (
         f"{LEFT_PARENTHESIS}{column_definition}{RIGHT_PARENTHESIS}"
@@ -877,8 +880,9 @@ def create_table_as_select_statement(
         }
     )
     return (
-        f"{CREATE}{OR + REPLACE if replace else EMPTY_STRING} {table_type.upper()} {TABLE}"
-        f"{IF + NOT + EXISTS if not replace and not error else EMPTY_STRING} "
+        f"{CREATE}{OR + REPLACE if replace else EMPTY_STRING}"
+        f" {(get_temp_type_for_object(use_scoped_temp_objects, is_generated) if table_type.lower() in TEMPORARY_STRING_SET else table_type).upper()} "
+        f"{TABLE}{IF + NOT + EXISTS if not replace and not error else EMPTY_STRING} "
         f"{table_name}{column_definition_sql}{cluster_by_clause}{options_statement}"
         f"{COPY_GRANTS if copy_grants else EMPTY_STRING}{comment_sql} {AS}{project_statement([], child)}"
     )

src/snowflake/snowpark/_internal/analyzer/snowflake_plan.py

@@ -929,6 +929,8 @@ def get_create_table_as_select_plan(child: SnowflakePlan, replace, error):
                     max_data_extension_time=max_data_extension_time,
                     change_tracking=change_tracking,
                     copy_grants=copy_grants,
+                    use_scoped_temp_objects=use_scoped_temp_objects,
+                    is_generated=is_generated,
                 ),
                 child,
                 source_plan,

tests/integ/compiler/test_query_generator.py

@@ -197,7 +197,7 @@ def test_table_create_from_large_query_breakdown(session, plan_source_generator)
 
     assert (
         queries[PlanQueryType.QUERIES][0].sql
-        == f" CREATE  TEMP  TABLE  {table_name}    AS  SELECT  *  FROM (select 1 as a, 2 as b)"
+        == f" CREATE  SCOPED TEMPORARY  TABLE  {table_name}    AS  SELECT  *  FROM (select 1 as a, 2 as b)"
     )
 
 

tests/integ/scala/test_snowflake_plan_suite.py

@@ -317,7 +317,7 @@ def test_create_scoped_temp_table(session):
             )
             .queries[0]
             .sql
-            == f" CREATE  TEMP  TABLE  {temp_table_name}    AS  SELECT  *  FROM ( SELECT  *  FROM ({table_name}))"
+            == f" CREATE  TEMPORARY  TABLE  {temp_table_name}    AS  SELECT  *  FROM ( SELECT  *  FROM ({table_name}))"
         )
         expected_sql = f' CREATE  TEMPORARY  TABLE  {temp_table_name}("NUM" BIGINT, "STR" STRING(8))'
         assert expected_sql in (
@@ -342,7 +342,9 @@ def test_create_scoped_temp_table(session):
             .queries[0]
             .sql
         )
-        expected_sql = f" CREATE  TEMPORARY  TABLE  {temp_table_name}    AS  SELECT"
+        expected_sql = (
+            f" CREATE  SCOPED TEMPORARY  TABLE  {temp_table_name}    AS  SELECT"
+        )
         assert expected_sql in (
             session._plan_builder.save_as_table(
                 table_name=[temp_table_name],

tests/integ/test_large_query_breakdown.py

@@ -99,7 +99,7 @@ def test_large_query_breakdown_with_cte_optimization(session):
     check_result_with_and_without_breakdown(session, df4)
 
     assert len(df4.queries["queries"]) == 2
-    assert df4.queries["queries"][0].startswith("CREATE  TEMP  TABLE")
+    assert df4.queries["queries"][0].startswith("CREATE  SCOPED TEMPORARY  TABLE")
     assert df4.queries["queries"][1].startswith("WITH SNOWPARK_TEMP_CTE_")
 
     assert len(df4.queries["post_actions"]) == 1
@@ -115,7 +115,7 @@ def test_save_as_table(session, large_query_df):
 
     assert len(history.queries) == 4
     assert history.queries[0].sql_text == "SELECT CURRENT_TRANSACTION()"
-    assert history.queries[1].sql_text.startswith("CREATE  TEMP  TABLE")
+    assert history.queries[1].sql_text.startswith("CREATE  SCOPED TEMPORARY  TABLE")
     assert history.queries[2].sql_text.startswith(
         f"CREATE  OR  REPLACE    TABLE  {table_name}"
     )
@@ -135,7 +135,7 @@ def test_update_delete_merge(session, large_query_df):
         t.update({"B": 0}, t.a == large_query_df.a, large_query_df)
     assert len(history.queries) == 4
     assert history.queries[0].sql_text == "SELECT CURRENT_TRANSACTION()"
-    assert history.queries[1].sql_text.startswith("CREATE  TEMP  TABLE")
+    assert history.queries[1].sql_text.startswith("CREATE  SCOPED TEMPORARY  TABLE")
     assert history.queries[2].sql_text.startswith(f"UPDATE {table_name}")
     assert history.queries[3].sql_text.startswith("DROP  TABLE  If  EXISTS")
 
@@ -144,7 +144,7 @@ def test_update_delete_merge(session, large_query_df):
         t.delete(t.a == large_query_df.a, large_query_df)
     assert len(history.queries) == 4
     assert history.queries[0].sql_text == "SELECT CURRENT_TRANSACTION()"
-    assert history.queries[1].sql_text.startswith("CREATE  TEMP  TABLE")
+    assert history.queries[1].sql_text.startswith("CREATE  SCOPED TEMPORARY  TABLE")
     assert history.queries[2].sql_text.startswith(f"DELETE  FROM {table_name} USING")
     assert history.queries[3].sql_text.startswith("DROP  TABLE  If  EXISTS")
 
@@ -157,7 +157,7 @@ def test_update_delete_merge(session, large_query_df):
         )
     assert len(history.queries) == 4
     assert history.queries[0].sql_text == "SELECT CURRENT_TRANSACTION()"
-    assert history.queries[1].sql_text.startswith("CREATE  TEMP  TABLE")
+    assert history.queries[1].sql_text.startswith("CREATE  SCOPED TEMPORARY  TABLE")
     assert history.queries[2].sql_text.startswith(f"MERGE  INTO {table_name} USING")
     assert history.queries[3].sql_text.startswith("DROP  TABLE  If  EXISTS")
 
@@ -176,7 +176,7 @@ def test_copy_into_location(session, large_query_df):
         )
     assert len(history.queries) == 4, history.queries
     assert history.queries[0].sql_text == "SELECT CURRENT_TRANSACTION()"
-    assert history.queries[1].sql_text.startswith("CREATE  TEMP  TABLE")
+    assert history.queries[1].sql_text.startswith("CREATE  SCOPED TEMPORARY  TABLE")
     assert history.queries[2].sql_text.startswith(f"COPY  INTO '{remote_file_path}'")
     assert history.queries[3].sql_text.startswith("DROP  TABLE  If  EXISTS")
 
@@ -215,7 +215,7 @@ def test_pivot_unpivot(session):
 
     plan_queries = final_df.queries
     assert len(plan_queries["queries"]) == 2
-    assert plan_queries["queries"][0].startswith("CREATE  TEMP  TABLE")
+    assert plan_queries["queries"][0].startswith("CREATE  SCOPED TEMPORARY  TABLE")
 
     assert len(plan_queries["post_actions"]) == 1
     assert plan_queries["post_actions"][0].startswith("DROP  TABLE  If  EXISTS")
@@ -239,7 +239,7 @@ def test_sort(session):
 
     plan_queries = final_df.queries
     assert len(plan_queries["queries"]) == 2
-    assert plan_queries["queries"][0].startswith("CREATE  TEMP  TABLE")
+    assert plan_queries["queries"][0].startswith("CREATE  SCOPED TEMPORARY  TABLE")
 
     assert len(plan_queries["post_actions"]) == 1
     assert plan_queries["post_actions"][0].startswith("DROP  TABLE  If  EXISTS")
@@ -283,7 +283,7 @@ def test_multiple_query_plan(session, large_query_df):
             "CREATE  OR  REPLACE  SCOPED TEMPORARY  TABLE"
         )
         assert plan_queries["queries"][1].startswith("INSERT  INTO")
-        assert plan_queries["queries"][2].startswith("CREATE  TEMP  TABLE")
+        assert plan_queries["queries"][2].startswith("CREATE  SCOPED TEMPORARY  TABLE")
 
         assert len(plan_queries["post_actions"]) == 2
         for query in plan_queries["post_actions"]:
@@ -349,7 +349,9 @@ def test_async_job_with_large_query_breakdown(session, large_query_df):
     result = job.result()
     assert result == large_query_df.collect()
     assert len(large_query_df.queries["queries"]) == 2
-    assert large_query_df.queries["queries"][0].startswith("CREATE  TEMP  TABLE")
+    assert large_query_df.queries["queries"][0].startswith(
+        "CREATE  SCOPED TEMPORARY  TABLE"
+    )
 
     assert len(large_query_df.queries["post_actions"]) == 1
     assert large_query_df.queries["post_actions"][0].startswith(
@@ -365,7 +367,9 @@ def test_complexity_bounds_affect_num_partitions(session, large_query_df):
     session._large_query_breakdown_enabled = True
     assert len(large_query_df.queries["queries"]) == 2
     assert len(large_query_df.queries["post_actions"]) == 1
-    assert large_query_df.queries["queries"][0].startswith("CREATE  TEMP  TABLE")
+    assert large_query_df.queries["queries"][0].startswith(
+        "CREATE  SCOPED TEMPORARY  TABLE"
+    )
     assert large_query_df.queries["post_actions"][0].startswith(
         "DROP  TABLE  If  EXISTS"
     )
@@ -374,8 +378,12 @@ def test_complexity_bounds_affect_num_partitions(session, large_query_df):
     session._large_query_breakdown_enabled = True
     assert len(large_query_df.queries["queries"]) == 3
     assert len(large_query_df.queries["post_actions"]) == 2
-    assert large_query_df.queries["queries"][0].startswith("CREATE  TEMP  TABLE")
-    assert large_query_df.queries["queries"][1].startswith("CREATE  TEMP  TABLE")
+    assert large_query_df.queries["queries"][0].startswith(
+        "CREATE  SCOPED TEMPORARY  TABLE"
+    )
+    assert large_query_df.queries["queries"][1].startswith(
+        "CREATE  SCOPED TEMPORARY  TABLE"
+    )
     assert large_query_df.queries["post_actions"][0].startswith(
         "DROP  TABLE  If  EXISTS"
     )