add ability to specify vertorized_scanner for write_pandas

Author: culpgrant

src/snowflake/connector/pandas_tools.py

@@ -231,6 +231,7 @@ def write_pandas(
     overwrite: bool = False,
     table_type: Literal["", "temp", "temporary", "transient"] = "",
     use_logical_type: bool | None = None,
+    use_vectorized_scanner: bool | None = None,
     **kwargs: Any,
 ) -> tuple[
     bool,
@@ -295,6 +296,10 @@ def write_pandas(
             Snowflake can interpret Parquet logical 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:
             https://docs.snowflake.com/en/sql-reference/sql/create-file-format
+        use_vectorized_scanner: Boolean that specifies to use a vectorized scanner for loading Parquet files.
+            Using the vectorized scanner can significantly reduce the latency for loading Parquet files. To enable
+            Vectorized scanning of Parquet files, set use_vectorized_scanner as True. Set to None to use Snowflakes default.
+            For more information, see: https://docs.snowflake.com/en/sql-reference/sql/copy-into-table#label-use-vectorized-scanner
 
 
     Returns:
@@ -533,6 +538,7 @@ def drop_object(name: str, object_type: str) -> None:
             f"COMPRESSION={compression_map[compression]}"
             f"{' BINARY_AS_TEXT=FALSE' if auto_create_table or overwrite else ''}"
             f"{sql_use_logical_type}"
+            f"{' USE_VECTORIZED_SCANNER=' + str(use_vectorized_scanner).upper() if use_vectorized_scanner is not None else ''}"
             f") "
             f"PURGE=TRUE ON_ERROR=?"
         )

test/integ/pandas/test_pandas_tools.py

@@ -482,6 +482,47 @@ def test_write_pandas_use_logical_type(
             cnx.execute_string(drop_sql)
 
 
+@pytest.mark.parametrize(
+    ("use_vectorized_scanner", "expected_file_format"),
+    [
+        (None, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto)"),
+        (True, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto USE_VECTORIZED_SCANNER=TRUE)"),
+        (False, "FILE_FORMAT=(TYPE=PARQUET COMPRESSION=auto USE_VECTORIZED_SCANNER=FALSE)"),
+    ],
+)
+def test_write_pandas_use_vectorized_scanner(
+    conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
+    use_vectorized_scanner: bool | None,
+    expected_file_format: str,
+):
+    """Test that use_vectorized_scanner is making correct arguments to the COPY INTO command in SQL."""
+    from snowflake.connector.cursor import SnowflakeCursor
+
+    table_name = random_string(5, "use_vectorized_scanner")
+
+    with conn_cnx() as cnx:
+        def mocked_execute(*args, **kwargs):
+            if len(args) >= 1 and args[0].startswith("COPY INTO"):
+                assert expected_file_format in args[0]
+            cur = SnowflakeCursor(cnx)
+            cur._result = iter([])
+            return cur
+
+        with mock.patch(
+            "snowflake.connector.cursor.SnowflakeCursor.execute",
+            side_effect=mocked_execute,
+        ) as m_execute:
+            success, nchunks, nrows, _ = write_pandas(
+                cnx,
+                sf_connector_version_df.get(),
+                table_name=table_name,
+                use_vectorized_scanner=use_vectorized_scanner,
+            )
+            assert m_execute.called and any(
+                map(lambda e: "COPY INTO" in str(e[0]), m_execute.call_args_list)
+            )
+
+
 def test_invalid_table_type_write_pandas(
     conn_cnx: Callable[..., Generator[SnowflakeConnection, None, None]],
 ):
@@ -534,6 +575,7 @@ def test_table_location_building(
 
         def mocked_execute(*args, **kwargs):
             if len(args) >= 1 and args[0].startswith("COPY INTO"):
+                print(kwargs)
                 assert kwargs["params"][0] == expected_location
             cur = SnowflakeCursor(cnx)
             cur._result = iter([])