fix: convert rows in batches for fetchone() and cursor iteration

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Release History
 
+# Unreleased
+- Fix: iterating a cursor (`for row in cursor`) and calling `fetchone()` in a loop are now about as fast as `fetchall()`. They previously converted each row through pandas on its own, costing around 1 ms per row, which made SQLAlchemy result iteration very slow. Rows are now converted `arraysize` at a time and served from an internal buffer.
+
 # 4.6.0 (2026-09-24)
 - Upgrade Databricks SQL Kernel to 1.1.0; the kernel dependency is now stable and no longer experimental.
 - Transparently auto-recover Thrift connections to Reyden / Real-Time warehouses: when a warehouse rejects the default Thrift protocol (SQLSTATE `KP001`), the session is re-opened on the kernel backend and the warehouse is remembered so later connections skip Thrift. Applies only when no backend was chosen explicitly.

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

@@ -32,7 +32,7 @@
 
 import logging
 from collections import deque
-from typing import Any, Deque, List, Optional, TYPE_CHECKING, cast
+from typing import Any, Deque, List, TYPE_CHECKING, cast
 
 import pyarrow
 
@@ -190,37 +190,35 @@ def _drain(self) -> pyarrow.Table:
     # ----- Arrow fetches -----
 
     def fetchall_arrow(self) -> pyarrow.Table:
-        return self._drain()
+        buffered = self._take_buffered_table()
+        table = self._drain()
+        if buffered is None:
+            return table
+        return pyarrow.concat_tables([buffered, table])
 
     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._take_buffered_table(size)
+        n_remaining = size if buffered is None else size - buffered.num_rows
+        self._ensure_buffered(n_remaining)
+        table = self._take_buffered(n_remaining)
+        if buffered is None:
+            return table
+        return pyarrow.concat_tables([buffered, table])
 
     # ----- 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
+    def _fetchmany_table(self, size: int) -> pyarrow.Table:
+        return self.fetchmany_arrow(size)
 
-    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)
+    def _fetchall_table(self) -> pyarrow.Table:
+        return self.fetchall_arrow()
 
-    def fetchall(self) -> List[Row]:
-        return self._convert_arrow_table(self._drain())
+    def _convert_table(self, table: pyarrow.Table) -> List[Row]:
+        return self._convert_arrow_table(table)
 
     def close(self) -> None:
         """Close the underlying kernel handle and notify the backend.

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

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, List, Optional, TYPE_CHECKING
+from typing import Any, List, Optional, TYPE_CHECKING, Union
 
 import logging
 
@@ -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,11 @@ 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)
+        buffered = self._take_buffered_table(size) or []
+        results = self.results.next_n_rows(size - len(buffered))
         self._next_row_index += len(results)
 
-        return results
+        return buffered + results
 
     def fetchall_json(self) -> List[List[str]]:
         """
@@ -173,10 +175,11 @@ def fetchall_json(self) -> List[List[str]]:
             Columnar table containing all remaining rows
         """
 
+        buffered = self._take_buffered_table() or []
         results = self.results.remaining_rows()
         self._next_row_index += len(results)
 
-        return results
+        return buffered + results
 
     def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """
@@ -196,71 +199,53 @@ 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)
         if isinstance(self.results, JsonQueue):
-            results = self._convert_json_to_arrow_table(results)
+            return self._convert_json_to_arrow_table(self.fetchmany_json(size))
 
+        buffered = self._take_buffered_table(size)
+        n_buffered = 0 if buffered is None else buffered.num_rows
+        results = self.results.next_n_rows(size - n_buffered)
         self._next_row_index += results.num_rows
 
-        return results
+        return self._prepend_buffered(buffered, results)
 
     def fetchall_arrow(self) -> "pyarrow.Table":
         """
         Fetch all remaining rows as an Arrow table.
         """
 
-        results = self.results.remaining_rows()
         if isinstance(self.results, JsonQueue):
-            results = self._convert_json_to_arrow_table(results)
+            return self._convert_json_to_arrow_table(self.fetchall_json())
 
+        buffered = self._take_buffered_table()
+        results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
 
-        return results
-
-    def fetchone(self) -> Optional[Row]:
-        """
-        Fetch the next row of a query result set, returning a single sequence,
-        or None when no more data is available.
+        return self._prepend_buffered(buffered, results)
 
-        Returns:
-            A single Row object or None if no more rows are available
-        """
+    @staticmethod
+    def _prepend_buffered(
+        buffered: Optional["pyarrow.Table"], results: "pyarrow.Table"
+    ) -> "pyarrow.Table":
+        if buffered is None:
+            return results
+        if results.num_rows == 0:
+            return buffered
+        return concat_table_chunks([buffered, results])
 
+    def _fetchmany_table(self, size: int) -> Union[List[List[str]], "pyarrow.Table"]:
         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
-
-    def fetchmany(self, size: int) -> List[Row]:
-        """
-        Fetch the next set of rows of a query result, returning a list of rows.
-
-        Args:
-            size: Number of rows to fetch (defaults to arraysize if None)
-
-        Returns:
-            List of Row objects
-
-        Raises:
-            ValueError: If size is negative
-        """
+            return self.fetchmany_json(size)
+        return self.fetchmany_arrow(size)
 
+    def _fetchall_table(self) -> Union[List[List[str]], "pyarrow.Table"]:
         if isinstance(self.results, JsonQueue):
-            return self._create_json_table(self.fetchmany_json(size))
-        else:
-            return self._convert_arrow_table(self.fetchmany_arrow(size))
-
-    def fetchall(self) -> List[Row]:
-        """
-        Fetch all remaining rows of a query result, returning them as a list of rows.
-
-        Returns:
-            List of Row objects containing all remaining rows
-        """
+            return self.fetchall_json()
+        return self.fetchall_arrow()
 
+    def _convert_table(
+        self, table: Union[List[List[str]], "pyarrow.Table"]
+    ) -> 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._create_json_table(table)
+        return self._convert_arrow_table(table)