Iterating a cursor or calling fetchone() in a loop is 100x+ slower than fetchall() (affects SQLAlchemy) (#981)

src/databricks/sql/backend/kernel/result_set.py

@@ -189,38 +189,77 @@ def _drain(self) -> pyarrow.Table:
 
     # ----- Arrow fetches -----
 
+    def _fetchmany_arrow_unbuffered(self, size: int) -> pyarrow.Table:
+        """Take up to ``size`` rows from the batch FIFO, bypassing the
+        fetchone() row buffer."""
+        if size == 0:
+            return pyarrow.Table.from_batches([], schema=self._schema)
+        self._ensure_buffered(size)
+        return self._take_buffered(size)
+
+    def _fetch_available(self, size: int) -> pyarrow.Table:
+        """Take up to ``size`` rows already buffered, for fetchone() refills.
+
+        Unlike ``_fetchmany_arrow_unbuffered`` this does not keep pulling
+        batches until ``size`` rows are collected: it only pulls from the
+        kernel while the batch FIFO is empty, so latency to the first row
+        matches the pre-buffering behaviour.
+        """
+        self._ensure_buffered(1)
+        return self._take_buffered(min(size, self._buffered_rows()))
+
+    @staticmethod
+    def _concat_buffered_arrow(
+        buffered: Optional[pyarrow.Table], rest: pyarrow.Table
+    ) -> pyarrow.Table:
+        """Prepend rows drained from the fetchone() buffer to ``rest``.
+
+        Both tables are built against the result set's ``_schema``, so a plain
+        ``pyarrow.concat_tables`` suffices (no ColumnTable / schema
+        promotion handling as in ``ThriftResultSet._prepend_buffered``).
+        """
+        if buffered is None:
+            return rest
+        if rest.num_rows == 0:
+            return buffered
+        return pyarrow.concat_tables([buffered, rest])
+
     def fetchall_arrow(self) -> pyarrow.Table:
-        return self._drain()
+        buffered = self._row_buffer_take_raw()
+        return self._concat_buffered_arrow(buffered, self._drain())
 
     def fetchmany_arrow(self, size: int) -> pyarrow.Table:
         if size < 0:
             raise ValueError(f"fetchmany_arrow size must be >= 0, got {size}")
-        if size == 0:
-            return pyarrow.Table.from_batches([], schema=self._schema)
-        self._ensure_buffered(size)
-        return self._take_buffered(size)
+        buffered = self._row_buffer_take_raw(size)
+        remaining = size - (buffered.num_rows if buffered is not None else 0)
+        if buffered is not None and remaining == 0:
+            return buffered
+        return self._concat_buffered_arrow(
+            buffered, self._fetchmany_arrow_unbuffered(remaining)
+        )
 
     # ----- Row fetches -----
 
     def fetchone(self) -> Optional[Row]:
-        self._ensure_buffered(1)
-        if self._buffered_rows() == 0:
-            return None
-        table = self._take_buffered(1)
-        rows = self._convert_arrow_table(table)
-        return rows[0] if rows else None
+        # Rows are converted ``arraysize`` at a time and served from the
+        # base-class row buffer; converting one row per call is very slow.
+        return self._fetchone_buffered(
+            self._fetch_available, self._convert_arrow_table
+        )
 
     def fetchmany(self, size: int) -> List[Row]:
         if size < 0:
             raise ValueError(f"fetchmany size must be >= 0, got {size}")
-        if size == 0:
-            return []
-        self._ensure_buffered(size)
-        table = self._take_buffered(size)
-        return self._convert_arrow_table(table)
+        return self._fetch_rows_with_buffer(
+            size,
+            lambda n: self._convert_arrow_table(self._fetchmany_arrow_unbuffered(n)),
+        )
 
     def fetchall(self) -> List[Row]:
-        return self._convert_arrow_table(self._drain())
+        return self._fetch_rows_with_buffer(
+            None, lambda _: self._convert_arrow_table(self._drain())
+        )
 
     def close(self) -> None:
         """Close the underlying kernel handle and notify the backend.
@@ -234,6 +273,10 @@ def close(self) -> None:
         connection-close would otherwise issue a kernel call into an
         already-disposed session.
         """
+        # Drop rows parked by fetchone() (and any lost-rows poison) so
+        # post-close fetches return nothing, matching the cleared batch FIFO below.
+        self._row_buffer_clear()
+        self._lost_rows_error = None
         if self._kernel_handle is None:
             return
         if not self.connection.open:

src/databricks/sql/backend/sea/result_set.py

@@ -19,6 +19,7 @@
 from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory
 from databricks.sql.backend.types import ExecuteResponse
 from databricks.sql.result_set import ResultSet
+from databricks.sql.utils import concat_table_chunks
 
 logger = logging.getLogger(__name__)
 
@@ -160,10 +161,7 @@ def fetchmany_json(self, size: int) -> List[List[str]]:
         if size < 0:
             raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
 
-        results = self.results.next_n_rows(size)
-        self._next_row_index += len(results)
-
-        return results
+        return self._fetch_raw_with_buffer(size)
 
     def fetchall_json(self) -> List[List[str]]:
         """
@@ -173,11 +171,48 @@ def fetchall_json(self) -> List[List[str]]:
             Columnar table containing all remaining rows
         """
 
-        results = self.results.remaining_rows()
-        self._next_row_index += len(results)
+        return self._fetch_raw_with_buffer(None)
+
+    @staticmethod
+    def _raw_num_rows(raw: Any) -> int:
+        return len(raw) if isinstance(raw, list) else raw.num_rows
 
+    def _fetch_raw_from_queue(self, size: Optional[int]) -> Any:
+        """Fetch ``size`` (all if None) raw rows from the queue, bypassing the row buffer."""
+        if size is None:
+            results = self.results.remaining_rows()
+        else:
+            results = self.results.next_n_rows(size)
+        self._next_row_index += self._raw_num_rows(results)
         return results
 
+    def _fetch_available_from_queue(self, size: int) -> Any:
+        """Fetch up to ``size`` rows the queue already holds, for Arrow fetchone() refills.
+
+        On CloudFetch, the request is capped at what the current chunk still holds
+        (matching Thrift) via the shared ``ResultSet._cap_to_current_chunk``, so a
+        refill is limited to one chunk rather than stitching several links into one
+        batch. It does not avoid waiting on the next download: reading to the end of
+        a chunk still makes ``next_n_rows`` advance to the next link in that call.
+        Only called while the fetchone() row buffer is empty.
+        """
+        return self._fetch_raw_from_queue(self._cap_to_current_chunk(size))
+
+    def _fetch_raw_with_buffer(self, size: Optional[int]) -> Any:
+        """Fetch ``size`` (all if None) raw rows, draining the fetchone() buffer first."""
+        buffered = self._row_buffer_take_raw(size)
+        if buffered is None:
+            return self._fetch_raw_from_queue(size)
+        remaining = None if size is None else size - self._raw_num_rows(buffered)
+        if remaining == 0:
+            return buffered
+        rest = self._fetch_raw_from_queue(remaining)
+        if self._raw_num_rows(rest) == 0:
+            return buffered
+        if isinstance(buffered, list):
+            return buffered + list(rest)
+        return concat_table_chunks([buffered, rest])
+
     def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """
         Fetch the next set of rows as an Arrow table.
@@ -196,25 +231,21 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         if size < 0:
             raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
 
-        results = self.results.next_n_rows(size)
+        results = self._fetch_raw_with_buffer(size)
         if isinstance(self.results, JsonQueue):
             results = self._convert_json_to_arrow_table(results)
 
-        self._next_row_index += results.num_rows
-
         return results
 
     def fetchall_arrow(self) -> "pyarrow.Table":
         """
         Fetch all remaining rows as an Arrow table.
         """
 
-        results = self.results.remaining_rows()
+        results = self._fetch_raw_with_buffer(None)
         if isinstance(self.results, JsonQueue):
             results = self._convert_json_to_arrow_table(results)
 
-        self._next_row_index += results.num_rows
-
         return results
 
     def fetchone(self) -> Optional[Row]:
@@ -227,11 +258,10 @@ def fetchone(self) -> Optional[Row]:
         """
 
         if isinstance(self.results, JsonQueue):
-            res = self._create_json_table(self.fetchmany_json(1))
-        else:
-            res = self._convert_arrow_table(self.fetchmany_arrow(1))
-
-        return res[0] if res else None
+            return self._fetchone_buffered(self.fetchmany_json, self._create_json_table)
+        return self._fetchone_buffered(
+            self._fetch_available_from_queue, self._convert_arrow_table
+        )
 
     def fetchmany(self, size: int) -> List[Row]:
         """
@@ -247,10 +277,15 @@ def fetchmany(self, size: int) -> List[Row]:
             ValueError: If size is negative
         """
 
+        if size < 0:
+            raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
         if isinstance(self.results, JsonQueue):
-            return self._create_json_table(self.fetchmany_json(size))
-        else:
-            return self._convert_arrow_table(self.fetchmany_arrow(size))
+            return self._fetch_rows_with_buffer(
+                size, lambda n: self._create_json_table(self.fetchmany_json(n))
+            )
+        return self._fetch_rows_with_buffer(
+            size, lambda n: self._convert_arrow_table(self.fetchmany_arrow(n))
+        )
 
     def fetchall(self) -> List[Row]:
         """
@@ -261,6 +296,9 @@ def fetchall(self) -> List[Row]:
         """
 
         if isinstance(self.results, JsonQueue):
-            return self._create_json_table(self.fetchall_json())
-        else:
-            return self._convert_arrow_table(self.fetchall_arrow())
+            return self._fetch_rows_with_buffer(
+                None, lambda _: self._create_json_table(self.fetchall_json())
+            )
+        return self._fetch_rows_with_buffer(
+            None, lambda _: self._convert_arrow_table(self.fetchall_arrow())
+        )

src/databricks/sql/backend/sea/utils/filters.py

@@ -174,8 +174,9 @@ def _filter_arrow_result_set(
             raise ValueError(f"Column index {column_index} is out of bounds")
         column_name = result_set.description[column_index][0]
 
-        # Get all remaining rows as Arrow table and filter it
-        arrow_table = result_set.results.remaining_rows()
+        # Get all remaining rows as Arrow table and filter it. Go through
+        # fetchall_arrow() so any rows buffered by fetchone() are included.
+        arrow_table = result_set.fetchall_arrow()
         filtered_table = ResultSetFilter._filter_arrow_table(
             arrow_table, column_name, allowed_values, case_sensitive
         )
@@ -220,8 +221,9 @@ def _filter_json_result_set(
         if column_index >= len(result_set.description):
             raise ValueError(f"Column index {column_index} is out of bounds")
 
-        # Extract rows
-        all_rows = result_set.results.remaining_rows()
+        # Extract rows. Go through fetchall_json() so any rows buffered by
+        # fetchone() are included.
+        all_rows = result_set.fetchall_json()
 
         # Convert allowed values if case-insensitive
         if not case_sensitive: