Skip to content

Commit 4d29402

Browse files
SNOW-1882338 Fix AST to_snowpark_pandas test in nop testing mode (#2863)
1. Which Jira issue is this PR addressing? Make sure that there is an accompanying issue to your PR. Fixes SNOW-1882338 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. - [x] I acknowledge that I have ensured my changes to be thread-safe. Follow the link for more information: [Thread-safe Developer Guidelines](https://github.com/snowflakedb/snowpark-python/blob/main/CONTRIBUTING.md#thread-safe-development) 3. Please describe how your code solves the related issue. Add logic to the no-op AST testing code to create a temp read-only table for to_snowpark_pandas. This way `to_snowpark_pandas` works with both `nop` mode and `local_testing_mode`.
1 parent e633e96 commit 4d29402

File tree

3 files changed

+34
-3
lines changed

3 files changed

+34
-3
lines changed

src/snowflake/snowpark/mock/_nop_connection.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,10 @@
2121
from snowflake.connector.cursor import ResultMetadata
2222
from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
2323
from snowflake.snowpark._internal.analyzer.expression import Attribute
24-
from snowflake.snowpark._internal.analyzer.snowflake_plan import SnowflakePlan
24+
from snowflake.snowpark._internal.analyzer.snowflake_plan import (
25+
SnowflakePlan,
26+
PlanQueryType,
27+
)
2528
from snowflake.snowpark._internal.analyzer.snowflake_plan_node import (
2629
LogicalPlan,
2730
SaveMode,
@@ -120,6 +123,28 @@ def execute(
120123
]:
121124
source_plan = plan.source_plan
122125

126+
if hasattr(source_plan, "execution_queries"):
127+
# If temp read-only table, explicitly create it.
128+
# This occurs when code such as to_snowpark_pandas is run where the Snowpark version of the table is
129+
# cloned and then read.
130+
from snowflake.snowpark.mock import TableEmulator
131+
132+
for plan_query_type, query in source_plan.execution_queries.items():
133+
if query:
134+
query_sql = query[0].sql
135+
if (
136+
plan_query_type == PlanQueryType.QUERIES
137+
and "TEMPORARY READ ONLY TABLE" in query_sql
138+
):
139+
temp_table_name = query_sql.split("TEMPORARY READ ONLY TABLE ")[
140+
1
141+
].split(" ")[0]
142+
self.entity_registry.write_table(
143+
temp_table_name,
144+
TableEmulator({"A": [1], "B": [1], "C": [1]}),
145+
SaveMode.IGNORE,
146+
)
147+
123148
if isinstance(source_plan, SnowflakeCreateTable):
124149
result = self.entity_registry.write_table(
125150
source_plan.table_name,

src/snowflake/snowpark/mock/_nop_plan.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from typing import Any, Dict, List, Optional
77

88
import snowflake.snowpark
9+
from snowflake.snowpark.mock import TableEmulator
910
from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
1011
from snowflake.snowpark._internal.analyzer.binary_plan_node import Join
1112
from snowflake.snowpark._internal.analyzer.expression import (
@@ -117,6 +118,9 @@ def resolve_attributes(
117118
pivot_attrs.extend(pivot_result_cols)
118119
attributes = pivot_attrs
119120

121+
elif isinstance(plan, TableEmulator):
122+
attributes = [Attribute(name, _NumericType(), False) for name in plan.columns]
123+
120124
elif isinstance(plan, TableUpdate):
121125
attributes = [
122126
Attribute(name, _NumericType(), False)

tests/ast/data/Dataframe.to_snowpark_pandas.test

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,10 @@ body {
3030
expr {
3131
sp_table {
3232
name {
33-
sp_table_name_flat {
34-
name: "table1"
33+
name {
34+
sp_name_flat {
35+
name: "table1"
36+
}
3537
}
3638
}
3739
src {

0 commit comments

Comments
 (0)