SNOW-2217588:add support for USE_VECTORIZED_SCANNER in write_pandas() (#3759)

Author: sfc-gh-yuwang

CHANGELOG.md

@@ -10,6 +10,7 @@
 - Added a new function `interval_year_month_from_parts` that allows users to easily create `YearMonthIntervalType` without using SQL.
 - Added support for `FileOperation.list` to list files in a stage with metadata.
 - Added support for `FileOperation.remove` to remove files in a stage.
+- Added support for parameter `use_vectorized_scanner` in function `Session.write_pandas()`.
 
 #### Bug Fixes
 

src/snowflake/snowpark/session.py

@@ -3190,6 +3190,7 @@ def write_pandas(
         overwrite: bool = False,
         table_type: Literal["", "temp", "temporary", "transient"] = "",
         use_logical_type: Optional[bool] = None,
+        use_vectorized_scanner: bool = False,
         _emit_ast: bool = True,
         **kwargs: Dict[str, Any],
     ) -> Table:
@@ -3235,6 +3236,8 @@ def write_pandas(
                 types during data loading. To enable Parquet logical types, set use_logical_type as True. Set to None to
                 use Snowflakes default. For more information, see:
                 `file format options: <https://docs.snowflake.com/en/sql-reference/sql/create-file-format#type-parquet>`_.
+            use_vectorized_scanner: Boolean that specifies whether to use a vectorized scanner for loading Parquet files. See details at
+                `copy options <https://docs.snowflake.com/en/sql-reference/sql/copy-into-table.html#copy-options-copyoptions>`_.
 
         Example::
 
@@ -3374,6 +3377,7 @@ def write_pandas(
                         auto_create_table=auto_create_table,
                         overwrite=overwrite,
                         table_type=table_type,
+                        use_vectorized_scanner=use_vectorized_scanner,
                         **kwargs,
                     )
         except ProgrammingError as pe:

tests/integ/test_pandas_to_df.py

@@ -8,6 +8,7 @@
 import logging
 import math
 from datetime import date, datetime, time, timedelta, timezone
+from unittest import mock
 
 import pytest
 
@@ -325,6 +326,41 @@ def test_write_pandas_with_table_type(session, table_type: str):
         Utils.drop_table(session, table_name)
 
 
+@pytest.mark.skipif(
+    "config.getoption('local_testing_mode', default=False)",
+    reason="SNOW-1421323: session.write_pandas is not supported in Local Testing.",
+)
+@pytest.mark.parametrize("use_vectorized_scanner", [False, True])
+def test_write_pandas_use_vectorized_scanner(session, use_vectorized_scanner):
+    from snowflake.connector.pandas_tools import write_pandas
+
+    original_func = write_pandas
+
+    def wrapper(*args, **kwargs):
+        return original_func(*args, **kwargs)
+
+    pd = PandasDF(
+        [
+            (1, 4.5, "t1"),
+            (2, 7.5, "t2"),
+            (3, 10.5, "t3"),
+        ],
+        columns=["id".upper(), "foot_size".upper(), "shoe_model".upper()],
+    )
+    tb_name = Utils.random_name_for_temp_object(TempObjectType.TABLE)
+    with mock.patch(
+        "snowflake.snowpark.session.write_pandas", side_effect=wrapper
+    ) as execute:
+        session.write_pandas(
+            pd,
+            tb_name,
+            auto_create_table=True,
+            use_vectorized_scanner=use_vectorized_scanner,
+        )
+        _, kwargs = execute.call_args
+        assert kwargs["use_vectorized_scanner"] == use_vectorized_scanner
+
+
 @pytest.mark.skipif(
     "config.getoption('local_testing_mode', default=False)",
     reason="SNOW-1421323: session.write_pandas is not supported in Local Testing.",