Skip to content

SNOW-2057867 refactor and fixes to make pandas write work for Python … #2304

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 11 commits into
base: zyao-SNOW-2057867-refactor-stage-bind-to-make-it-work-for-sprocs
Choose a base branch
from
+34 −31
Open

SNOW-2057867 refactor and fixes to make pandas write work for Python … #2304

Show file tree
Hide file tree
Changes from 6 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
2722162
SNOW-2057867 refactor and fixes to make pandas write work for Python …
sfc-gh-zyao Apr 28, 2025
2af775b
revert the change to file_transfer_agent.py, because we want to follo…
sfc-gh-zyao Apr 28, 2025
f725b8f
Merge branch 'zyao-SNOW-2057867-refactor-stage-bind-to-make-it-work-f…
sfc-gh-zyao Apr 29, 2025
65ccd7c
Merge branch 'zyao-SNOW-2057867-refactor-stage-bind-to-make-it-work-f…
sfc-gh-zyao Apr 29, 2025
0d7ccc7
Merge branch 'zyao-SNOW-2057867-refactor-stage-bind-to-make-it-work-f…
sfc-gh-zyao Apr 29, 2025
b3a2c32
Merge branch 'zyao-SNOW-2057867-refactor-stage-bind-to-make-it-work-f…
sfc-gh-zyao Apr 29, 2025
ed8f6bd
use single quotes to surround the local path, to keep it consistent w…
sfc-gh-zyao Apr 29, 2025
03d10a2
Merge branch 'zyao-SNOW-2057867-refactor-stage-bind-to-make-it-work-f…
sfc-gh-zyao Apr 30, 2025
ad96279
Merge branch 'zyao-SNOW-2057867-refactor-stage-bind-to-make-it-work-f…
sfc-gh-zyao Apr 30, 2025
30ce232
put "source_compression": "auto_detect" in the options
sfc-gh-zyao Apr 30, 2025
e841eb4
fix a pandas_tools bug, where we don't use snowpark naming pattern fo…
sfc-gh-zyao Apr 30, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 15 additions & 16 deletions src/snowflake/connector/pandas_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -383,6 +383,10 @@ def write_pandas(
"Unsupported table type. Expected table types: temp/temporary, transient"
)

if table_type.lower() in ["temp", "temporary"]:
# Add scoped keyword when applicable.
table_type = get_temp_type_for_object(_use_scoped_temp_object).lower()

if chunk_size is None:
chunk_size = len(df)

Expand Down Expand Up @@ -438,22 +442,13 @@ def write_pandas(
# Dump chunk into parquet file
chunk.to_parquet(chunk_path, compression=compression, **kwargs)
# Upload parquet file
upload_sql = (
"PUT /* Python:snowflake.connector.pandas_tools.write_pandas() */ "
"'file://{path}' ? PARALLEL={parallel}"
).format(
path=chunk_path.replace("\\", "\\\\").replace("'", "\\'"),
parallel=parallel,
)
params = ("@" + stage_location,)
logger.debug(f"uploading files with '{upload_sql}', params: %s", params)
cursor.execute(
upload_sql,
_is_internal=True,
_force_qmark_paramstyle=True,
params=params,
num_statements=1,
path = chunk_path.replace("\\", "\\\\").replace("'", "\\'")
cursor._upload(
local_file_name="file://" + path,
stage_location="@" + stage_location,
options={"parallel": parallel},
)

# Remove chunk file
os.remove(chunk_path)

Expand Down Expand Up @@ -516,7 +511,11 @@ def drop_object(name: str, object_type: str) -> None:
target_table_location = build_location_helper(
database,
schema,
random_string() if (overwrite and auto_create_table) else table_name,
(
random_name_for_temp_object(TempObjectType.TABLE)
if (overwrite and auto_create_table)
else table_name
),
quote_identifiers,
)

Expand Down
21 changes: 17 additions & 4 deletions test/integ/pandas/test_pandas_tools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Any, Callable, Generator
from unittest import mock
from unittest.mock import MagicMock

import numpy.random
import pytest
Expand Down Expand Up @@ -543,7 +544,10 @@ def mocked_execute(*args, **kwargs):
with mock.patch(
"snowflake.connector.cursor.SnowflakeCursor.execute",
side_effect=mocked_execute,
) as m_execute:
) as m_execute, mock.patch(
"snowflake.connector.cursor.SnowflakeCursor._upload",
side_effect=MagicMock(),
) as _:
success, nchunks, nrows, _ = write_pandas(
cnx,
sf_connector_version_df.get(),
Expand Down Expand Up @@ -591,7 +595,10 @@ def mocked_execute(*args, **kwargs):
with mock.patch(
"snowflake.connector.cursor.SnowflakeCursor.execute",
side_effect=mocked_execute,
) as m_execute:
) as m_execute, mock.patch(
"snowflake.connector.cursor.SnowflakeCursor._upload",
side_effect=MagicMock(),
) as _:
success, nchunks, nrows, _ = write_pandas(
cnx,
sf_connector_version_df.get(),
Expand Down Expand Up @@ -643,7 +650,10 @@ def mocked_execute(*args, **kwargs):
with mock.patch(
"snowflake.connector.cursor.SnowflakeCursor.execute",
side_effect=mocked_execute,
) as m_execute:
) as m_execute, mock.patch(
"snowflake.connector.cursor.SnowflakeCursor._upload",
side_effect=MagicMock(),
) as _:
cnx._update_parameters({"PYTHON_SNOWPARK_USE_SCOPED_TEMP_OBJECTS": True})
success, nchunks, nrows, _ = write_pandas(
cnx,
Expand Down Expand Up @@ -701,7 +711,10 @@ def mocked_execute(*args, **kwargs):
with mock.patch(
"snowflake.connector.cursor.SnowflakeCursor.execute",
side_effect=mocked_execute,
) as m_execute:
) as m_execute, mock.patch(
"snowflake.connector.cursor.SnowflakeCursor._upload",
side_effect=MagicMock(),
) as _:
success, nchunks, nrows, _ = write_pandas(
cnx,
sf_connector_version_df.get(),
Expand Down
Loading