SNOW-806291: Fix Snowpark much slower than Spark Connector on Databricks (for collect() and toPandas()) (#1946)

Author: sfc-gh-yuwang

DESCRIPTION.md

@@ -8,6 +8,10 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 # Release Notes
 
+- v3.12.0(TBD)
+  - Optimized `to_pandas()` performance by fully parallel downloading logic.
+
+
 - v3.11.0(June 17,2024)
   - Added support for `token_file_path` connection parameter to read an OAuth token from a file when connecting to Snowflake.
   - Added support for `debug_arrow_chunk` connection parameter to allow debugging raw arrow data in case of arrow data parsing failure.

src/snowflake/connector/result_set.py

@@ -6,7 +6,7 @@
 
 import inspect
 from collections import deque
-from concurrent.futures import Future
+from concurrent.futures import ALL_COMPLETED, Future, wait
 from concurrent.futures.thread import ThreadPoolExecutor
 from logging import getLogger
 from typing import (
@@ -61,45 +61,64 @@ def result_set_iterator(
     Just like ``ResultBatch`` iterator, this might yield an ``Exception`` to allow users
     to continue iterating through the rest of the ``ResultBatch``.
     """
-
-    with ThreadPoolExecutor(prefetch_thread_num) as pool:
-        # Fill up window
-
-        logger.debug("beginning to schedule result batch downloads")
-
-        for _ in range(min(prefetch_thread_num, len(unfetched_batches))):
-            logger.debug(
-                f"queuing download of result batch id: {unfetched_batches[0].id}"
-            )
-            unconsumed_batches.append(
-                pool.submit(unfetched_batches.popleft().create_iter, **kw)
-            )
-
-        yield from first_batch_iter
-
-        i = 1
-        while unconsumed_batches:
-            logger.debug(f"user requesting to consume result batch {i}")
-
-            # Submit the next un-fetched batch to the pool
-            if unfetched_batches:
+    is_fetch_all = kw.pop("is_fetch_all", False)
+    if is_fetch_all:
+        with ThreadPoolExecutor(prefetch_thread_num) 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)
                 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()
+                logger.debug(f"user began consuming result batch {i}")
+                i += 1
+        final()
+    else:
+        with ThreadPoolExecutor(prefetch_thread_num) as pool:
+            # Fill up window
+
+            logger.debug("beginning to schedule result batch downloads")
+
+            for _ in range(min(prefetch_thread_num, len(unfetched_batches))):
+                logger.debug(
+                    f"queuing download of result batch id: {unfetched_batches[0].id}"
+                )
+                unconsumed_batches.append(
+                    pool.submit(unfetched_batches.popleft().create_iter, **kw)
+                )
 
-            future = unconsumed_batches.popleft()
+            yield from first_batch_iter
 
-            # this will raise an exception if one has occurred
-            batch_iterator = future.result()
+            i = 1
+            while unconsumed_batches:
+                logger.debug(f"user requesting to consume result batch {i}")
 
-            logger.debug(f"user began consuming result batch {i}")
-            yield from batch_iterator
-            logger.debug(f"user finished consuming result batch {i}")
+                # Submit the next un-fetched batch to the pool
+                if unfetched_batches:
+                    logger.debug(
+                        f"queuing download of result batch id: {unfetched_batches[0].id}"
+                    )
+                    future = pool.submit(unfetched_batches.popleft().create_iter, **kw)
+                    unconsumed_batches.append(future)
 
-            i += 1
-    final()
+                future = unconsumed_batches.popleft()
+
+                # this will raise an exception if one has occurred
+                batch_iterator = future.result()
+
+                logger.debug(f"user began consuming result batch {i}")
+                yield from batch_iterator
+                logger.debug(f"user finished consuming result batch {i}")
+
+                i += 1
+        final()
 
 
 class ResultSet(Iterable[list]):
@@ -202,7 +221,7 @@ def _fetch_pandas_all(self, **kwargs) -> DataFrame:
         """Fetches a single Pandas dataframe."""
         concat_args = list(inspect.signature(pandas.concat).parameters)
         concat_kwargs = {k: kwargs.pop(k) for k in dict(kwargs) if k in concat_args}
-        dataframes = list(self._fetch_pandas_batches(**kwargs))
+        dataframes = list(self._fetch_pandas_batches(is_fetch_all=True, **kwargs))
         if dataframes:
             return pandas.concat(
                 dataframes,
@@ -238,6 +257,9 @@ def _create_iter(
         This function is a helper function to ``__iter__`` and it was introduced for the
         cases where we need to propagate some values to later ``_download`` calls.
         """
+        # pop is_fetch_all and pass it to result_set_iterator
+        is_fetch_all = kwargs.pop("is_fetch_all", False)
+
         # add connection so that result batches can use sessions
         kwargs["connection"] = self._cursor.connection
 
@@ -257,6 +279,7 @@ def _create_iter(
             unfetched_batches,
             self._finish_iterating,
             self.prefetch_thread_num,
+            is_fetch_all=is_fetch_all,
             **kwargs,
         )