SNOW-2062305 process pool batch fetcher (#2365)

Author: sfc-gh-mmishchenko

DESCRIPTION.md

@@ -11,7 +11,8 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Bumped numpy dependency from <2.1.0 to <=2.2.4
   - Added Windows support for Python 3.13.
   - Add `bulk_upload_chunks` parameter to `write_pandas` function. Setting this parameter to True changes the behaviour of write_pandas function to first write all the data chunks to the local disk and then perform the wildcard upload of the chunks folder to the stage. In default behaviour the chunks are being saved, uploaded and deleted one by one.
-  - Added support for new authentication mechanism PAT with external session ID
+  - Added support for new authentication mechanism PAT with external session ID.
+  - Added `client_fetch_use_mp` parameter that enables multiprocessed fetching of result batches.
 
 - v3.15.1(May 20, 2025)
   - Added basic arrow support for Interval types.

src/snowflake/connector/connection.py

@@ -229,6 +229,7 @@ def _get_private_bytes_from_file(
     ),  # snowflake
     "client_prefetch_threads": (4, int),  # snowflake
     "client_fetch_threads": (None, (type(None), int)),
+    "client_fetch_use_mp": (False, bool),
     "numpy": (False, bool),  # snowflake
     "ocsp_response_cache_filename": (None, (type(None), str)),  # snowflake internal
     "converter_class": (DefaultConverterClass(), SnowflakeConverter),
@@ -428,7 +429,9 @@ class SnowflakeConnection:
             See the backoff_policies module for details and implementation examples.
         client_session_keep_alive_heartbeat_frequency: Heartbeat frequency to keep connection alive in seconds.
         client_prefetch_threads: Number of threads to download the result set.
-        client_fetch_threads: Number of threads to fetch staged query results.
+        client_fetch_threads: Number of threads (or processes) to fetch staged query results.
+            If not specified, reuses client_prefetch_threads value.
+        client_fetch_use_mp: Enables multiprocessing for fetching query results in parallel.
         rest: Snowflake REST API object. Internal use only. Maybe removed in a later release.
         application: Application name to communicate with Snowflake as. By default, this is "PythonConnector".
         errorhandler: Handler used with errors. By default, an exception will be raised on error.
@@ -701,6 +704,10 @@ def client_fetch_threads(self, value: None | int) -> None:
             value = min(max(1, value), MAX_CLIENT_FETCH_THREADS)
         self._client_fetch_threads = value
 
+    @property
+    def client_fetch_use_mp(self) -> bool:
+        return self._client_fetch_use_mp
+
     @property
     def rest(self) -> SnowflakeRestful | None:
         return self._rest

src/snowflake/connector/cursor.py

@@ -1212,6 +1212,7 @@ def _init_result_and_meta(self, data: dict[Any, Any]) -> None:
             result_chunks,
             self._connection.client_fetch_threads
             or self._connection.client_prefetch_threads,
+            self._connection.client_fetch_use_mp,
         )
         self._rownumber = -1
         self._result_state = ResultState.VALID

src/snowflake/connector/result_batch.py

@@ -8,6 +8,8 @@
 from logging import getLogger
 from typing import TYPE_CHECKING, Any, Callable, Iterator, NamedTuple, Sequence
 
+from typing_extensions import Self
+
 from .arrow_context import ArrowConverterContext
 from .backoff_policies import exponential_backoff
 from .compat import OK, UNAUTHORIZED, urlparse
@@ -413,6 +415,14 @@ def to_pandas(self) -> DataFrame:
     def to_arrow(self) -> Table:
         raise NotImplementedError()
 
+    @abc.abstractmethod
+    def populate_data(
+        self, connection: SnowflakeConnection | None = None, **kwargs
+    ) -> Self:
+        """Downloads the data that the ``ResultBatch`` is pointing at and populates it into self._data.
+        Returns the instance itself."""
+        raise NotImplementedError()
+
 
 class JSONResultBatch(ResultBatch):
     def __init__(
@@ -538,11 +548,9 @@ def _parse(
     def __repr__(self) -> str:
         return f"JSONResultChunk({self.id})"
 
-    def create_iter(
+    def _fetch_data(
         self, connection: SnowflakeConnection | None = None, **kwargs
-    ) -> Iterator[dict | Exception] | Iterator[tuple | Exception]:
-        if self._local:
-            return iter(self._data)
+    ) -> list[dict | Exception] | list[tuple | Exception]:
         response = self._download(connection=connection)
         # Load data to a intermediate form
         logger.debug(f"started loading result batch id: {self.id}")
@@ -554,7 +562,20 @@ def create_iter(
         with TimerContextManager() as parse_metric:
             parsed_data = self._parse(downloaded_data)
         self._metrics[DownloadMetrics.parse.value] = parse_metric.get_timing_millis()
-        return iter(parsed_data)
+        return parsed_data
+
+    def populate_data(
+        self, connection: SnowflakeConnection | None = None, **kwargs
+    ) -> Self:
+        self._data = self._fetch_data(connection=connection, **kwargs)
+        return self
+
+    def create_iter(
+        self, connection: SnowflakeConnection | None = None, **kwargs
+    ) -> Iterator[dict | Exception] | Iterator[tuple | Exception]:
+        if self._local:
+            return iter(self._data)
+        return iter(self._fetch_data(connection=connection, **kwargs))
 
     def _arrow_fetching_error(self):
         return NotSupportedError(
@@ -613,7 +634,10 @@ def _load(
         )
 
     def _from_data(
-        self, data: str, iter_unit: IterUnit, check_error_on_every_column: bool = True
+        self,
+        data: str | bytes,
+        iter_unit: IterUnit,
+        check_error_on_every_column: bool = True,
     ) -> Iterator[dict | Exception] | Iterator[tuple | Exception]:
         """Creates a ``PyArrowIterator`` files from a str.
 
@@ -623,8 +647,11 @@ def _from_data(
         if len(data) == 0:
             return iter([])
 
+        if isinstance(data, str):
+            data = b64decode(data)
+
         return _create_nanoarrow_iterator(
-            b64decode(data),
+            data,
             self._context,
             self._use_dict_result,
             self._numpy,
@@ -751,3 +778,9 @@ def create_iter(
                 return self._get_arrow_iter(connection=connection)
         else:
             return self._create_iter(iter_unit=iter_unit, connection=connection)
+
+    def populate_data(
+        self, connection: SnowflakeConnection | None = None, **kwargs
+    ) -> Self:
+        self._data = self._download(connection=connection).content
+        return self

src/snowflake/connector/result_set.py

@@ -2,7 +2,7 @@
 
 import inspect
 from collections import deque
-from concurrent.futures import ALL_COMPLETED, Future, wait
+from concurrent.futures import ALL_COMPLETED, Future, ProcessPoolExecutor, wait
 from concurrent.futures.thread import ThreadPoolExecutor
 from logging import getLogger
 from typing import (
@@ -44,6 +44,7 @@ def result_set_iterator(
     unfetched_batches: Deque[ResultBatch],
     final: Callable[[], None],
     prefetch_thread_num: int,
+    use_mp: bool,
     **kw: Any,
 ) -> Iterator[dict | Exception] | Iterator[tuple | Exception] | Iterator[Table]:
     """Creates an iterator over some other iterators.
@@ -58,26 +59,52 @@ def result_set_iterator(
     to continue iterating through the rest of the ``ResultBatch``.
     """
     is_fetch_all = kw.pop("is_fetch_all", False)
+
+    if use_mp:
+
+        def create_pool_executor() -> ProcessPoolExecutor:
+            return ProcessPoolExecutor(prefetch_thread_num)
+
+        def create_fetch_task(batch: ResultBatch):
+            return batch.populate_data
+
+        def get_fetch_result(future_result: ResultBatch):
+            return future_result.create_iter(**kw)
+
+        kw["connection"] = None
+    else:
+
+        def create_pool_executor() -> ThreadPoolExecutor:
+            return ThreadPoolExecutor(prefetch_thread_num)
+
+        def create_fetch_task(batch: ResultBatch):
+            return batch.create_iter
+
+        def get_fetch_result(future_result: Iterator):
+            return future_result
+
     if is_fetch_all:
-        with ThreadPoolExecutor(prefetch_thread_num) as pool:
+        with create_pool_executor() as pool:
             logger.debug("beginning to schedule result batch downloads")
             yield from first_batch_iter
             while unfetched_batches:
                 logger.debug(
                     f"queuing download of result batch id: {unfetched_batches[0].id}"
                 )
-                future = pool.submit(unfetched_batches.popleft().create_iter, **kw)
+                future = pool.submit(
+                    create_fetch_task(unfetched_batches.popleft()), **kw
+                )
                 unconsumed_batches.append(future)
             _, _ = wait(unconsumed_batches, return_when=ALL_COMPLETED)
             i = 1
             while unconsumed_batches:
                 logger.debug(f"user began consuming result batch {i}")
-                yield from unconsumed_batches.popleft().result()
+                yield from get_fetch_result(unconsumed_batches.popleft().result())
                 logger.debug(f"user began consuming result batch {i}")
                 i += 1
         final()
     else:
-        with ThreadPoolExecutor(prefetch_thread_num) as pool:
+        with create_pool_executor() as pool:
             # Fill up window
 
             logger.debug("beginning to schedule result batch downloads")
@@ -87,7 +114,7 @@ def result_set_iterator(
                     f"queuing download of result batch id: {unfetched_batches[0].id}"
                 )
                 unconsumed_batches.append(
-                    pool.submit(unfetched_batches.popleft().create_iter, **kw)
+                    pool.submit(create_fetch_task(unfetched_batches.popleft()), **kw)
                 )
 
             yield from first_batch_iter
@@ -101,13 +128,15 @@ def result_set_iterator(
                     logger.debug(
                         f"queuing download of result batch id: {unfetched_batches[0].id}"
                     )
-                    future = pool.submit(unfetched_batches.popleft().create_iter, **kw)
+                    future = pool.submit(
+                        create_fetch_task(unfetched_batches.popleft()), **kw
+                    )
                     unconsumed_batches.append(future)
 
                 future = unconsumed_batches.popleft()
 
                 # this will raise an exception if one has occurred
-                batch_iterator = future.result()
+                batch_iterator = get_fetch_result(future.result())
 
                 logger.debug(f"user began consuming result batch {i}")
                 yield from batch_iterator
@@ -136,10 +165,12 @@ def __init__(
         cursor: SnowflakeCursor,
         result_chunks: list[JSONResultBatch] | list[ArrowResultBatch],
         prefetch_thread_num: int,
+        use_mp: bool,
     ) -> None:
         self.batches = result_chunks
         self._cursor = cursor
         self.prefetch_thread_num = prefetch_thread_num
+        self._use_mp = use_mp
 
     def _report_metrics(self) -> None:
         """Report all metrics totalled up.
@@ -276,6 +307,7 @@ def _create_iter(
             self._finish_iterating,
             self.prefetch_thread_num,
             is_fetch_all=is_fetch_all,
+            use_mp=self._use_mp,
             **kwargs,
         )
 

test/integ/pandas_it/test_arrow_pandas.py

@@ -1285,9 +1285,10 @@ def test_to_arrow_datatypes(enable_structured_types, conn_cnx):
                         cur.execute(f"alter session unset {param}")
 
 
-def test_simple_arrow_fetch(conn_cnx):
+@pytest.mark.parametrize("client_fetch_use_mp", [False, True])
+def test_simple_arrow_fetch(conn_cnx, client_fetch_use_mp):
     rowcount = 250_000
-    with conn_cnx() as cnx:
+    with conn_cnx(client_fetch_use_mp=client_fetch_use_mp) as cnx:
         with cnx.cursor() as cur:
             cur.execute(SQL_ENABLE_ARROW)
             cur.execute(
@@ -1316,8 +1317,9 @@ def test_simple_arrow_fetch(conn_cnx):
             assert lo == rowcount
 
 
-def test_arrow_zero_rows(conn_cnx):
-    with conn_cnx() as cnx:
+@pytest.mark.parametrize("client_fetch_use_mp", [False, True])
+def test_arrow_zero_rows(conn_cnx, client_fetch_use_mp):
+    with conn_cnx(client_fetch_use_mp=client_fetch_use_mp) as cnx:
         with cnx.cursor() as cur:
             cur.execute(SQL_ENABLE_ARROW)
             cur.execute("select 1::NUMBER(38,0) limit 0")

test/integ/test_cursor.py

@@ -1581,11 +1581,13 @@ def test__log_telemetry_job_data(conn_cnx, caplog):
         ("arrow", ArrowResultBatch),
     ),
 )
+@pytest.mark.parametrize("client_fetch_use_mp", [False, True])
 def test_resultbatch(
     conn_cnx,
     result_format,
     expected_chunk_type,
     capture_sf_telemetry,
+    client_fetch_use_mp,
 ):
     """This test checks the following things:
     1. After executing a query can we pickle the result batches
@@ -1598,7 +1600,8 @@ def test_resultbatch(
     with conn_cnx(
         session_parameters={
             "python_connector_query_result_format": result_format,
-        }
+        },
+        client_fetch_use_mp=client_fetch_use_mp,
     ) as con:
         with capture_sf_telemetry.patch_connection(con) as telemetry_data:
             with con.cursor() as cur: