From a0b3c9f8deec2ad9c51e15b1aad86d187311faae Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <3815206+peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 21:49:38 +0000
Subject: [PATCH 01/29] Iterating a cursor or calling fetchone() in a loop is
 100x+ slower than fetchall() (affects SQLAlchemy) (#981)

Signed-off-by: peco-engineer-bot[bot] <3815206+peco-engineer-bot[bot]@users.noreply.github.com>
---
 .../sql/backend/kernel/result_set.py          |  58 ++++--
 src/databricks/sql/backend/sea/result_set.py  |  67 +++++--
 src/databricks/sql/result_set.py              | 182 +++++++++++++++---
 tests/e2e/test_driver.py                      | 106 ++++++++++
 tests/e2e/test_kernel_backend.py              |  65 +++++++
 tests/unit/test_fetches.py                    |  77 ++++++++
 tests/unit/test_sea_result_set.py             |  77 ++++++++
 7 files changed, 564 insertions(+), 68 deletions(-)

diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index 15f78bf5f..800b9deea 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -189,38 +189,62 @@ 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)
+
+    @staticmethod
+    def _prepend_buffered(
+        buffered: Optional[pyarrow.Table], rest: pyarrow.Table
+    ) -> pyarrow.Table:
+        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._prepend_buffered(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._prepend_buffered(
+            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._fetchmany_arrow_unbuffered, 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)
+        buffered = self._row_buffer_take_rows(size)
+        remaining = size - len(buffered)
+        if remaining == 0:
+            return buffered
+        rest = self._convert_arrow_table(self._fetchmany_arrow_unbuffered(remaining))
+        return buffered + rest if buffered else rest
 
     def fetchall(self) -> List[Row]:
-        return self._convert_arrow_table(self._drain())
+        buffered = self._row_buffer_take_rows()
+        rest = self._convert_arrow_table(self._drain())
+        return buffered + rest if buffered else rest
 
     def close(self) -> None:
         """Close the underlying kernel handle and notify the backend.
diff --git a/src/databricks/sql/backend/sea/result_set.py b/src/databricks/sql/backend/sea/result_set.py
index 17838ed81..babcd9c8b 100644
--- a/src/databricks/sql/backend/sea/result_set.py
+++ b/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,36 @@ 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_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,12 +219,10 @@ 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":
@@ -209,12 +230,10 @@ 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 +246,8 @@ 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.fetchmany_arrow, self._convert_arrow_table)
 
     def fetchmany(self, size: int) -> List[Row]:
         """
@@ -247,10 +263,17 @@ 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")
+        buffered = self._row_buffer_take_rows(size)
+        remaining = size - len(buffered)
+        if buffered and remaining == 0:
+            return buffered
         if isinstance(self.results, JsonQueue):
-            return self._create_json_table(self.fetchmany_json(size))
+            rest = self._create_json_table(self.fetchmany_json(remaining))
         else:
-            return self._convert_arrow_table(self.fetchmany_arrow(size))
+            rest = self._convert_arrow_table(self.fetchmany_arrow(remaining))
+        return buffered + rest if buffered else rest
 
     def fetchall(self) -> List[Row]:
         """
@@ -260,7 +283,9 @@ def fetchall(self) -> List[Row]:
             List of Row objects containing all remaining rows
         """
 
+        buffered = self._row_buffer_take_rows()
         if isinstance(self.results, JsonQueue):
-            return self._create_json_table(self.fetchall_json())
+            rest = self._create_json_table(self.fetchall_json())
         else:
-            return self._convert_arrow_table(self.fetchall_arrow())
+            rest = self._convert_arrow_table(self.fetchall_arrow())
+        return buffered + rest if buffered else rest
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index f24a07505..adec9410f 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import List, Optional, TYPE_CHECKING, Tuple
+from typing import Any, Callable, List, Optional, TYPE_CHECKING, Tuple
 
 import logging
 import pandas
@@ -85,6 +85,15 @@ def __init__(
         self._arrow_schema_bytes = arrow_schema_bytes
         # Affected-row count for DML; None for SELECT / unreported.
         self.num_modified_rows = num_modified_rows
+        # Small row buffer that serves fetchone()/iteration. Rows are fetched and
+        # converted ``arraysize`` at a time instead of one at a time (converting a
+        # single row costs a full pandas round-trip). ``_row_buffer_raw`` holds the
+        # un-converted chunk (Arrow table, ColumnTable or JSON row list) aligned with
+        # ``_row_buffer_rows`` so the Arrow/columnar fetch methods can drain it too.
+        # Every fetch method drains this buffer first, preserving row order.
+        self._row_buffer_raw: Any = None
+        self._row_buffer_rows: List[Row] = []
+        self._row_buffer_pos = 0
 
     def __iter__(self):
         while True:
@@ -136,6 +145,74 @@ def _convert_arrow_table(self, table):
     def rownumber(self):
         return self._next_row_index
 
+    # ----- fetchone() row buffer helpers -----
+
+    def _row_buffer_count(self) -> int:
+        return len(self._row_buffer_rows) - self._row_buffer_pos
+
+    def _row_buffer_clear(self) -> None:
+        self._row_buffer_raw = None
+        self._row_buffer_rows = []
+        self._row_buffer_pos = 0
+
+    def _row_buffer_advance(self, n: Optional[int]) -> Tuple[int, int]:
+        """Advance the buffer by up to ``n`` rows (all if None); return (start, length)."""
+        available = self._row_buffer_count()
+        length = available if n is None else max(0, min(n, available))
+        start = self._row_buffer_pos
+        self._row_buffer_pos += length
+        # Buffered rows only count as consumed once they are handed to the caller.
+        self._next_row_index += length
+        return start, length
+
+    def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
+        """Pop up to ``n`` (all if None) converted rows from the buffer."""
+        if self._row_buffer_count() == 0:
+            return []
+        rows = self._row_buffer_rows
+        start, length = self._row_buffer_advance(n)
+        if self._row_buffer_count() == 0:
+            self._row_buffer_clear()
+        return rows[start : start + length]
+
+    def _row_buffer_take_raw(self, n: Optional[int] = None) -> Any:
+        """Pop up to ``n`` (all if None) un-converted rows from the buffer.
+
+        Returns None when the buffer is empty, otherwise a slice of the raw chunk
+        (same type as what the queue returned: Arrow table, ColumnTable or list).
+        """
+        if self._row_buffer_count() == 0:
+            return None
+        raw = self._row_buffer_raw
+        start, length = self._row_buffer_advance(n)
+        if self._row_buffer_count() == 0:
+            self._row_buffer_clear()
+        if isinstance(raw, list):
+            return raw[start : start + length]
+        return raw.slice(start, length)
+
+    def _fetchone_buffered(
+        self, fetch_raw: Callable[[int], Any], convert: Callable[[Any], List[Row]]
+    ) -> Optional[Row]:
+        """Serve fetchone() from the row buffer, refilling ``arraysize`` rows at a time.
+
+        ``_next_row_index`` keeps counting rows handed to the caller: rows parked in
+        the buffer are subtracted on refill and added back as they are taken. The
+        queue is only read again once the buffer is empty, so at that point the
+        index equals the queue position again (Thrift relies on that for
+        ``expected_row_start_offset``).
+        """
+        if self._row_buffer_count() == 0:
+            raw = fetch_raw(max(self.arraysize, 1))
+            rows = convert(raw)
+            if not rows:
+                return None
+            self._row_buffer_raw = raw
+            self._row_buffer_rows = rows
+            self._row_buffer_pos = 0
+            self._next_row_index -= len(rows)
+        return self._row_buffer_take_rows(1)[0]
+
     @property
     def is_staging_operation(self) -> bool:
         """Whether this result set represents a staging operation."""
@@ -302,12 +379,45 @@ def _convert_columnar_table(self, table):
 
         return result
 
+    def _prepend_buffered(self, buffered, rest):
+        """Concatenate rows drained from the fetchone() buffer with freshly fetched ones."""
+        if buffered is None:
+            return rest
+        if rest is None or rest.num_rows == 0:
+            return buffered
+        return concat_table_chunks([buffered, rest])
+
     def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         """
         Fetch the next set of rows of a query result, returning a PyArrow table.
 
         An empty sequence is returned when no more rows are available.
         """
+        if size < 0:
+            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+        buffered = self._row_buffer_take_raw(size)
+        remaining = size - (buffered.num_rows if buffered is not None else 0)
+        rest = self._fetchmany_arrow_from_queue(remaining) if remaining > 0 else None
+        if buffered is None and rest is None:
+            rest = self._fetchmany_arrow_from_queue(0)
+        return self._prepend_buffered(buffered, rest)
+
+    def fetchmany_columnar(self, size: int):
+        """
+        Fetch the next set of rows of a query result, returning a Columnar Table.
+        An empty sequence is returned when no more rows are available.
+        """
+        if size < 0:
+            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+        buffered = self._row_buffer_take_raw(size)
+        remaining = size - (buffered.num_rows if buffered is not None else 0)
+        rest = self._fetchmany_columnar_from_queue(remaining) if remaining > 0 else None
+        if buffered is None and rest is None:
+            rest = self._fetchmany_columnar_from_queue(0)
+        return self._prepend_buffered(buffered, rest)
+
+    def _fetchmany_arrow_from_queue(self, size: int) -> "pyarrow.Table":
+        """Fetch up to ``size`` rows from the results queue, bypassing the row buffer."""
         if size < 0:
             raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
 
@@ -343,11 +453,8 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
             partial_result_chunks.append(zero_row_table)
         return concat_table_chunks(partial_result_chunks)
 
-    def fetchmany_columnar(self, size: int):
-        """
-        Fetch the next set of rows of a query result, returning a Columnar Table.
-        An empty sequence is returned when no more rows are available.
-        """
+    def _fetchmany_columnar_from_queue(self, size: int):
+        """Fetch up to ``size`` rows from the results queue, bypassing the row buffer."""
         if size < 0:
             raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
 
@@ -370,6 +477,22 @@ def fetchmany_columnar(self, size: int):
 
     def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
+        buffered = self._row_buffer_take_raw()
+        result_table = self._prepend_buffered(buffered, self._fetchall_arrow_from_queue())
+        # If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table
+        # Valid only for metadata commands result set
+        if isinstance(result_table, ColumnTable) and pyarrow:
+            data = {
+                name: col
+                for name, col in zip(
+                    result_table.column_names, result_table.column_table
+                )
+            }
+            return pyarrow.Table.from_pydict(data)
+        return result_table
+
+    def _fetchall_arrow_from_queue(self):
+        """Drain the results queue (bypassing the row buffer) without type conversion."""
         # Hold 0-row chunks aside instead of appending them to ``partial_result_chunks``.
         # CloudFetchQueue may return a placeholder empty table whose schema does not
         # match the real downloaded chunks; concatenating it would corrupt the result.
@@ -394,21 +517,15 @@ def fetchall_arrow(self) -> "pyarrow.Table":
         if not partial_result_chunks:
             partial_result_chunks.append(zero_row_table)
 
-        result_table = concat_table_chunks(partial_result_chunks)
-        # If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table
-        # Valid only for metadata commands result set
-        if isinstance(result_table, ColumnTable) and pyarrow:
-            data = {
-                name: col
-                for name, col in zip(
-                    result_table.column_names, result_table.column_table
-                )
-            }
-            return pyarrow.Table.from_pydict(data)
-        return result_table
+        return concat_table_chunks(partial_result_chunks)
 
     def fetchall_columnar(self):
         """Fetch all (remaining) rows of a query result, returning them as a Columnar table."""
+        buffered = self._row_buffer_take_raw()
+        return self._prepend_buffered(buffered, self._fetchall_columnar_from_queue())
+
+    def _fetchall_columnar_from_queue(self):
+        """Drain the results queue as a Columnar table, bypassing the row buffer."""
         results = self.results.remaining_rows()
         self._next_row_index += results.num_rows
         partial_result_chunks = [results]
@@ -426,23 +543,21 @@ def fetchone(self) -> Optional[Row]:
         or None when no more data is available.
         """
         if isinstance(self.results, ColumnQueue):
-            res = self._convert_columnar_table(self.fetchmany_columnar(1))
-        else:
-            res = self._convert_arrow_table(self.fetchmany_arrow(1))
-
-        if len(res) > 0:
-            return res[0]
-        else:
-            return None
+            return self._fetchone_buffered(
+                self.fetchmany_columnar, self._convert_columnar_table
+            )
+        return self._fetchone_buffered(self.fetchmany_arrow, self._convert_arrow_table)
 
     def fetchall(self) -> List[Row]:
         """
         Fetch all (remaining) rows of a query result, returning them as a list of rows.
         """
+        buffered = self._row_buffer_take_rows()
         if isinstance(self.results, ColumnQueue):
-            return self._convert_columnar_table(self.fetchall_columnar())
+            rest = self._convert_columnar_table(self.fetchall_columnar())
         else:
-            return self._convert_arrow_table(self.fetchall_arrow())
+            rest = self._convert_arrow_table(self.fetchall_arrow())
+        return buffered + rest if buffered else rest
 
     def fetchmany(self, size: int) -> List[Row]:
         """
@@ -450,10 +565,17 @@ def fetchmany(self, size: int) -> List[Row]:
 
         An empty sequence is returned when no more rows are available.
         """
+        if size < 0:
+            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+        buffered = self._row_buffer_take_rows(size)
+        remaining = size - len(buffered)
+        if buffered and remaining == 0:
+            return buffered
         if isinstance(self.results, ColumnQueue):
-            return self._convert_columnar_table(self.fetchmany_columnar(size))
+            rest = self._convert_columnar_table(self.fetchmany_columnar(remaining))
         else:
-            return self._convert_arrow_table(self.fetchmany_arrow(size))
+            rest = self._convert_arrow_table(self.fetchmany_arrow(remaining))
+        return buffered + rest if buffered else rest
 
     @staticmethod
     def _get_schema_description(table_schema_message):
diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py
index 319b7de5e..ea99de595 100644
--- a/tests/e2e/test_driver.py
+++ b/tests/e2e/test_driver.py
@@ -942,6 +942,112 @@ def test_iterator_api(self, extra_params):
                 for j in range(len(row)):
                     assert row[j] == expected_results[i][j]
 
+    @contextmanager
+    def _count_row_conversions(self):
+        """Count every call that turns a fetched batch into ``Row`` objects
+        (Arrow, Thrift columnar and SEA JSON paths) — issue #981."""
+        from unittest import mock
+        from databricks.sql.result_set import ResultSet, ThriftResultSet
+        from databricks.sql.backend.sea.result_set import SeaResultSet
+
+        counter = {"calls": 0}
+        targets = [
+            (ResultSet, "_convert_arrow_table"),
+            (ThriftResultSet, "_convert_columnar_table"),
+            (SeaResultSet, "_create_json_table"),
+        ]
+        patches = []
+        for cls, name in targets:
+            orig = getattr(cls, name)
+
+            def _wrapped(self_, *args, _orig=orig, **kwargs):
+                counter["calls"] += 1
+                return _orig(self_, *args, **kwargs)
+
+            patches.append(mock.patch.object(cls, name, _wrapped))
+        for p in patches:
+            p.start()
+        try:
+            yield counter
+        finally:
+            for p in patches:
+                p.stop()
+
+    _ROW_BUFFER_PARAMS = [
+        {},
+        {
+            "use_sea": True,
+            "use_cloud_fetch": False,
+            "enable_query_result_lz4_compression": False,
+        },
+        {
+            "use_sea": True,
+        },
+    ]
+    _ROW_BUFFER_N = 5000
+    _ROW_BUFFER_QUERY = "SELECT id, CAST(id AS STRING) AS s FROM range({})"
+
+    def _assert_rows_and_bound(self, rows, calls, arraysize):
+        n = self._ROW_BUFFER_N
+        assert len(rows) == n
+        assert [r[0] for r in rows] == list(range(n))
+        assert all(r[1] == str(r[0]) for r in rows)
+        bound = -(-n // arraysize) + 2
+        assert calls <= bound, (
+            "row conversion called {} times for {} rows (arraysize={}); "
+            "expected at most {} (issue #981)".format(calls, n, arraysize, bound)
+        )
+
+    @pytest.mark.parametrize("extra_params", _ROW_BUFFER_PARAMS)
+    def test_iteration_converts_rows_in_batches(self, extra_params):
+        with self.cursor(extra_params) as cursor:
+            cursor.execute(self._ROW_BUFFER_QUERY.format(self._ROW_BUFFER_N))
+            with self._count_row_conversions() as counter:
+                rows = [row for row in cursor]
+            self._assert_rows_and_bound(rows, counter["calls"], cursor.arraysize)
+
+    @pytest.mark.parametrize("extra_params", _ROW_BUFFER_PARAMS)
+    def test_fetchone_loop_converts_rows_in_batches(self, extra_params):
+        with self.cursor(extra_params) as cursor:
+            cursor.execute(self._ROW_BUFFER_QUERY.format(self._ROW_BUFFER_N))
+            rows = []
+            with self._count_row_conversions() as counter:
+                while True:
+                    row = cursor.fetchone()
+                    if row is None:
+                        break
+                    rows.append(row)
+            self._assert_rows_and_bound(rows, counter["calls"], cursor.arraysize)
+
+    @pytest.mark.parametrize("extra_params", _ROW_BUFFER_PARAMS)
+    def test_fetchone_mixed_with_other_fetches_keeps_order(self, extra_params):
+        with self.cursor(extra_params) as cursor:
+            cursor.execute("SELECT * FROM range(10)")
+            assert cursor.fetchone()[0] == 0
+            assert cursor.fetchone()[0] == 1
+            assert cursor.rownumber == 2
+            assert [r[0] for r in cursor.fetchmany(3)] == [2, 3, 4]
+            assert cursor.rownumber == 5
+            assert cursor.fetchone()[0] == 5
+            assert cursor.rownumber == 6
+            table = cursor.fetchmany_arrow(2)
+            assert table.column(0).to_pylist() == [6, 7]
+            assert cursor.rownumber == 8
+            assert [r[0] for r in cursor.fetchall()] == [8, 9]
+            assert cursor.rownumber == 10
+            assert cursor.fetchone() is None
+
+    @pytest.mark.parametrize("extra_params", _ROW_BUFFER_PARAMS)
+    def test_fetchone_then_fetchall_arrow_keeps_order(self, extra_params):
+        with self.cursor(extra_params) as cursor:
+            cursor.execute("SELECT * FROM range(10)")
+            assert cursor.fetchone()[0] == 0
+            assert cursor.rownumber == 1
+            table = cursor.fetchall_arrow()
+            assert table.column(0).to_pylist() == list(range(1, 10))
+            assert cursor.rownumber == 10
+            assert cursor.fetchone() is None
+
     def test_temp_view_fetch(self):
         with self.cursor({}) as cursor:
             query = "create temporary view f as select * from range(10)"
diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py
index ae0301fb6..2b67f7df2 100644
--- a/tests/e2e/test_kernel_backend.py
+++ b/tests/e2e/test_kernel_backend.py
@@ -183,6 +183,71 @@ def test_fetchall_arrow(conn):
         assert table.column_names == ["a", "b"]
 
 
+def _count_arrow_conversions():
+    """Patch ``ResultSet._convert_arrow_table`` to count calls (issue #981)."""
+    from unittest import mock
+    from databricks.sql.result_set import ResultSet
+
+    orig = ResultSet._convert_arrow_table
+    counter = {"calls": 0}
+
+    def _wrapped(self_, table):
+        counter["calls"] += 1
+        return orig(self_, table)
+
+    return counter, mock.patch.object(ResultSet, "_convert_arrow_table", _wrapped)
+
+
+_ROW_BUFFER_N = 5000
+_ROW_BUFFER_ARRAYSIZE = 1000
+
+
+def _assert_rows_and_bound(rows, calls):
+    assert len(rows) == _ROW_BUFFER_N
+    assert [r[0] for r in rows] == list(range(_ROW_BUFFER_N))
+    assert all(r[1] == str(r[0]) for r in rows)
+    bound = -(-_ROW_BUFFER_N // _ROW_BUFFER_ARRAYSIZE) + 2
+    assert calls <= bound, (
+        f"_convert_arrow_table called {calls} times for {_ROW_BUFFER_N} rows; "
+        f"expected at most {bound} (issue #981)"
+    )
+
+
+def test_iteration_converts_rows_in_batches(conn):
+    with conn.cursor(arraysize=_ROW_BUFFER_ARRAYSIZE) as cur:
+        cur.execute(f"SELECT id, CAST(id AS STRING) AS s FROM range({_ROW_BUFFER_N})")
+        counter, patch = _count_arrow_conversions()
+        with patch:
+            rows = [row for row in cur]
+        _assert_rows_and_bound(rows, counter["calls"])
+
+
+def test_fetchone_loop_converts_rows_in_batches(conn):
+    with conn.cursor(arraysize=_ROW_BUFFER_ARRAYSIZE) as cur:
+        cur.execute(f"SELECT id, CAST(id AS STRING) AS s FROM range({_ROW_BUFFER_N})")
+        counter, patch = _count_arrow_conversions()
+        rows = []
+        with patch:
+            while True:
+                row = cur.fetchone()
+                if row is None:
+                    break
+                rows.append(row)
+        _assert_rows_and_bound(rows, counter["calls"])
+
+
+def test_fetchone_mixed_with_other_fetches_keeps_order(conn):
+    with conn.cursor() as cur:
+        cur.execute("SELECT * FROM range(10)")
+        assert cur.fetchone()[0] == 0
+        assert cur.fetchone()[0] == 1
+        assert [r[0] for r in cur.fetchmany(3)] == [2, 3, 4]
+        assert cur.fetchone()[0] == 5
+        assert cur.fetchmany_arrow(2).column(0).to_pylist() == [6, 7]
+        assert [r[0] for r in cur.fetchall()] == [8, 9]
+        assert cur.fetchone() is None
+
+
 @pytest.mark.parametrize("row_limit", [None, 0, 1, 5])
 def test_cursor_row_limit(conn, row_limit):
     with conn.cursor(row_limit=row_limit) as cur:
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index e3a963355..f613c4ac7 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -391,5 +391,82 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self):
         self.assertEqual(result.num_rows, 0)
 
 
+class ColumnarFetchoneBufferTests(unittest.TestCase):
+    """Issue #981: fetchone/iteration on the Thrift ColumnQueue path must convert
+    rows in batches (not one row at a time) while keeping order and rownumber."""
+
+    @staticmethod
+    def make_columnar_result_set(n_rows, arraysize):
+        from databricks.sql.utils import ColumnQueue, ColumnTable
+
+        table = ColumnTable([list(range(n_rows))], ["col0"])
+        mock_thrift_backend = Mock(spec=ThriftDatabricksClient)
+        mock_thrift_backend.fetch_results.return_value = (ColumnQueue(table), False, 0)
+        return ThriftResultSet(
+            connection=Mock(),
+            execute_response=ExecuteResponse(
+                command_id=None,
+                status=None,
+                has_been_closed_server_side=True,
+                description=[("col0", "integer", None, None, None, None, None)],
+                lz4_compressed=False,
+                is_staging_operation=False,
+            ),
+            thrift_client=mock_thrift_backend,
+            arraysize=arraysize,
+        )
+
+    @staticmethod
+    def _count_conversions():
+        from unittest.mock import patch
+
+        return patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=ThriftResultSet._convert_columnar_table,
+        )
+
+    def test_iteration_converts_in_batches(self):
+        rs = self.make_columnar_result_set(100, arraysize=10)
+        with self._count_conversions() as m:
+            rows = [r[0] for r in rs]
+        self.assertEqual(rows, list(range(100)))
+        self.assertLessEqual(m.call_count, 100 // 10 + 2)
+
+    def test_fetchone_loop_converts_in_batches(self):
+        rs = self.make_columnar_result_set(100, arraysize=10)
+        rows = []
+        with self._count_conversions() as m:
+            while True:
+                row = rs.fetchone()
+                if row is None:
+                    break
+                rows.append(row[0])
+        self.assertEqual(rows, list(range(100)))
+        self.assertLessEqual(m.call_count, 100 // 10 + 2)
+
+    def test_mixed_fetches_keep_order_and_rownumber(self):
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        self.assertEqual(rs.fetchone()[0], 0)
+        self.assertEqual(rs.fetchone()[0], 1)
+        self.assertEqual(rs.rownumber, 2)
+        self.assertEqual([r[0] for r in rs.fetchmany(3)], [2, 3, 4])
+        self.assertEqual(rs.rownumber, 5)
+        self.assertEqual(rs.fetchone()[0], 5)
+        self.assertEqual(rs.rownumber, 6)
+        self.assertEqual([r[0] for r in rs.fetchall()], [6, 7, 8, 9])
+        self.assertEqual(rs.rownumber, 10)
+        self.assertIsNone(rs.fetchone())
+
+    def test_fetchone_then_fetchall_arrow_keeps_order(self):
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        self.assertEqual(rs.fetchone()[0], 0)
+        self.assertEqual(rs.rownumber, 1)
+        table = rs.fetchall_arrow()
+        self.assertEqual(table.column(0).to_pylist(), list(range(1, 10)))
+        self.assertEqual(rs.rownumber, 10)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py
index 1c3e3b5b4..b004085a5 100644
--- a/tests/unit/test_sea_result_set.py
+++ b/tests/unit/test_sea_result_set.py
@@ -482,6 +482,83 @@ def test_iteration(self, result_set_with_data, sample_data):
         assert rows[0].col2 == 1
         assert rows[0].col3 is True
 
+    def _make_json_result_set(
+        self, mock_connection, mock_sea_client, execute_response, n_rows, arraysize
+    ):
+        data = [["v{}".format(i), str(i), "true"] for i in range(n_rows)]
+        result_data = ResultData(data=data, external_links=None, row_count=n_rows)
+        with patch(
+            "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue",
+            return_value=JsonQueue(data),
+        ):
+            return SeaResultSet(
+                connection=mock_connection,
+                execute_response=execute_response,
+                sea_client=mock_sea_client,
+                result_data=result_data,
+                manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY),
+                buffer_size_bytes=1000,
+                arraysize=arraysize,
+            )
+
+    @pytest.mark.parametrize("mode", ["iterate", "fetchone"])
+    def test_json_fetchone_converts_in_batches(
+        self, mock_connection, mock_sea_client, execute_response, mode
+    ):
+        """Issue #981: per-row fetchone must not convert one row at a time."""
+        rs = self._make_json_result_set(
+            mock_connection, mock_sea_client, execute_response, 100, arraysize=10
+        )
+        with patch.object(
+            SeaResultSet,
+            "_create_json_table",
+            autospec=True,
+            side_effect=SeaResultSet._create_json_table,
+        ) as m:
+            if mode == "iterate":
+                rows = list(rs)
+            else:
+                rows = []
+                while True:
+                    row = rs.fetchone()
+                    if row is None:
+                        break
+                    rows.append(row)
+        assert [r.col2 for r in rows] == list(range(100))
+        assert m.call_count <= 100 // 10 + 2
+
+    def test_json_mixed_fetches_keep_order_and_rownumber(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        rs = self._make_json_result_set(
+            mock_connection, mock_sea_client, execute_response, 10, arraysize=4
+        )
+        assert rs.fetchone().col2 == 0
+        assert rs.fetchone().col2 == 1
+        assert rs.rownumber == 2
+        assert [r.col2 for r in rs.fetchmany(3)] == [2, 3, 4]
+        assert rs.rownumber == 5
+        assert rs.fetchone().col2 == 5
+        assert rs.rownumber == 6
+        table = rs.fetchmany_arrow(2)
+        assert table.column(1).to_pylist() == [6, 7]
+        assert rs.rownumber == 8
+        assert [r.col2 for r in rs.fetchall()] == [8, 9]
+        assert rs.rownumber == 10
+        assert rs.fetchone() is None
+
+    def test_json_fetchone_then_fetchall_arrow_keeps_order(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        rs = self._make_json_result_set(
+            mock_connection, mock_sea_client, execute_response, 10, arraysize=4
+        )
+        assert rs.fetchone().col2 == 0
+        assert rs.rownumber == 1
+        table = rs.fetchall_arrow()
+        assert table.column(1).to_pylist() == list(range(1, 10))
+        assert rs.rownumber == 10
+
     def test_is_staging_operation(
         self, mock_connection, mock_sea_client, execute_response
     ):

From 57fb8a99dea2cf5cf0a89dca96bd2a285105821a Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 21:58:22 +0000
Subject: [PATCH 02/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4418300468 at src/databricks/sql/result_set.py:206

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 13 +++++++++++--
 tests/unit/test_fetches.py       | 11 +++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index adec9410f..cdfd1a0a4 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -27,6 +27,12 @@
 
 logger = logging.getLogger(__name__)
 
+# Upper bound on how many rows a fetchone() refill pulls and converts at once. The
+# refill uses ``arraysize`` but is capped here so the first fetchone() on a cursor
+# with a large arraysize (default 100000) doesn't materialise a full page just to
+# return one row.
+_FETCHONE_BUFFER_MAX_ROWS = 1000
+
 
 class ResultSet(ABC):
     """
@@ -194,7 +200,10 @@ def _row_buffer_take_raw(self, n: Optional[int] = None) -> Any:
     def _fetchone_buffered(
         self, fetch_raw: Callable[[int], Any], convert: Callable[[Any], List[Row]]
     ) -> Optional[Row]:
-        """Serve fetchone() from the row buffer, refilling ``arraysize`` rows at a time.
+        """Serve fetchone() from the row buffer, refilling in small batches.
+
+        Each refill fetches ``arraysize`` rows, capped at
+        ``_FETCHONE_BUFFER_MAX_ROWS``.
 
         ``_next_row_index`` keeps counting rows handed to the caller: rows parked in
         the buffer are subtracted on refill and added back as they are taken. The
@@ -203,7 +212,7 @@ def _fetchone_buffered(
         ``expected_row_start_offset``).
         """
         if self._row_buffer_count() == 0:
-            raw = fetch_raw(max(self.arraysize, 1))
+            raw = fetch_raw(max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1))
             rows = convert(raw)
             if not rows:
                 return None
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index f613c4ac7..3e3b2feb4 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -446,6 +446,17 @@ def test_fetchone_loop_converts_in_batches(self):
         self.assertEqual(rows, list(range(100)))
         self.assertLessEqual(m.call_count, 100 // 10 + 2)
 
+    def test_fetchone_refill_is_capped_for_large_arraysize(self):
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        n_rows = _FETCHONE_BUFFER_MAX_ROWS * 3
+        rs = self.make_columnar_result_set(n_rows, arraysize=100000)
+        self.assertEqual(rs.fetchone()[0], 0)
+        self.assertEqual(rs._row_buffer_count(), _FETCHONE_BUFFER_MAX_ROWS - 1)
+        self.assertEqual(rs.rownumber, 1)
+        self.assertEqual([r[0] for r in rs.fetchall()], list(range(1, n_rows)))
+        self.assertEqual(rs.rownumber, n_rows)
+
     def test_mixed_fetches_keep_order_and_rownumber(self):
         rs = self.make_columnar_result_set(10, arraysize=4)
         self.assertEqual(rs.fetchone()[0], 0)

From 2199868100f0bc9361836e85566ac97a6e16f20d Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:06:51 +0000
Subject: [PATCH 03/29] ai: apply changes for #1007 (3 review threads)

Addresses:
  - #4418360785 at src/databricks/sql/backend/kernel/result_set.py:228
  - #4418360790 at src/databricks/sql/result_set.py:558
  - #4418360798 at src/databricks/sql/result_set.py:219

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py     | 32 +++++++++++--
 tests/unit/test_fetches.py           | 66 ++++++++++++++++++++++++++-
 tests/unit/test_kernel_result_set.py | 67 +++++++++++++++++++++++++++-
 3 files changed, 158 insertions(+), 7 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index cdfd1a0a4..fd0042983 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -202,8 +202,10 @@ def _fetchone_buffered(
     ) -> Optional[Row]:
         """Serve fetchone() from the row buffer, refilling in small batches.
 
-        Each refill fetches ``arraysize`` rows, capped at
-        ``_FETCHONE_BUFFER_MAX_ROWS``.
+        Each refill asks ``fetch_raw`` for up to ``arraysize`` rows, capped at
+        ``_FETCHONE_BUFFER_MAX_ROWS``. ``fetch_raw`` may return fewer (Thrift only
+        serves what its current queue holds, refilling it at most until one row is
+        available, so latency to the first row is unchanged).
 
         ``_next_row_index`` keeps counting rows handed to the caller: rows parked in
         the buffer are subtracted on refill and added back as they are taken. The
@@ -484,6 +486,26 @@ def _fetchmany_columnar_from_queue(self, size: int):
 
         return concat_table_chunks(partial_result_chunks)
 
+    def _fetch_available_from_queue(self, size: int):
+        """Fetch up to ``size`` rows the queue already holds, for fetchone() refills.
+
+        Unlike ``_fetchmany_*_from_queue``, this does not keep calling the server
+        until ``size`` rows are collected: it only refills the queue while it is
+        empty, so fetchone() makes no more round-trips than it needs to return one
+        row (server pages can be smaller than ``size`` when they hit the byte limit).
+        Only called while the fetchone() row buffer is empty.
+        """
+        results = self.results.next_n_rows(size)
+        while (
+            results.num_rows == 0
+            and not self.has_been_closed_server_side
+            and self.has_more_rows
+        ):
+            self._fill_results_buffer()
+            results = self.results.next_n_rows(size)
+        self._next_row_index += results.num_rows
+        return results
+
     def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
         buffered = self._row_buffer_take_raw()
@@ -553,9 +575,11 @@ def fetchone(self) -> Optional[Row]:
         """
         if isinstance(self.results, ColumnQueue):
             return self._fetchone_buffered(
-                self.fetchmany_columnar, self._convert_columnar_table
+                self._fetch_available_from_queue, self._convert_columnar_table
             )
-        return self._fetchone_buffered(self.fetchmany_arrow, self._convert_arrow_table)
+        return self._fetchone_buffered(
+            self._fetch_available_from_queue, self._convert_arrow_table
+        )
 
     def fetchall(self) -> List[Row]:
         """
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 3e3b2feb4..8726c59ac 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -95,7 +95,8 @@ def make_dummy_result_set_from_initial_results(initial_results):
         return rs
 
     @staticmethod
-    def make_dummy_result_set_from_batch_list(batch_list):
+    def make_dummy_result_set_from_batch_list(batch_list, offsets=None):
+        """``offsets``, if given, records ``expected_row_start_offset`` per refill."""
         batch_index = 0
 
         def fetch_results(
@@ -110,6 +111,8 @@ def fetch_results(
             chunk_id=0,
         ):
             nonlocal batch_index
+            if offsets is not None:
+                offsets.append(expected_row_start_offset)
             results = FetchTests.make_arrow_queue(batch_list[batch_index])
             batch_index += 1
 
@@ -328,6 +331,67 @@ def test_fetchone_without_initial_results(self):
         dummy_result_set = self.make_dummy_result_set_from_batch_list(batch_list_2)
         self.assertEqual(dummy_result_set.fetchone(), None)
 
+    def test_mixed_fetches_across_pages_keep_order_and_refill_offsets(self):
+        # Issue #981: rows parked in the fetchone() buffer must not skew
+        # ``_next_row_index``, which Thrift sends as ``expected_row_start_offset``
+        # on every refill. Pages are 3 rows; fetchone() refills 2 rows at a time,
+        # so buffer refills and server page boundaries interleave.
+        batch_list = [[[0], [1], [2]], [[3], [4], [5]], [[6], [7], [8]], [[9]]]
+        offsets = []
+        rs = self.make_dummy_result_set_from_batch_list(batch_list, offsets=offsets)
+        rs.arraysize = 2
+        self.assertEqual(offsets, [0])
+
+        self.assertSequenceEqual(rs.fetchone(), [0])
+        self.assertEqual(rs.rownumber, 1)
+
+        # Drains the buffered row, the rest of page 0, then refills page 1.
+        table = rs.fetchmany_arrow(3)
+        self.assertEqual(table.column(0).to_pylist(), [1, 2, 3])
+        self.assertEqual(rs.rownumber, 4)
+        self.assertEqual(offsets, [0, 3])
+
+        self.assertSequenceEqual(rs.fetchone(), [4])
+        self.assertSequenceEqual(rs.fetchone(), [5])
+        self.assertEqual(rs.rownumber, 6)
+        self.assertEqual(offsets, [0, 3])
+
+        # Buffer empty and page 1 exhausted: the refill crosses into page 2.
+        self.assertSequenceEqual(rs.fetchone(), [6])
+        self.assertEqual(rs.rownumber, 7)
+        self.assertEqual(offsets, [0, 3, 6])
+
+        # Drains the buffered row [7], the rest of page 2, then page 3.
+        self.assertEqualRowValues(rs.fetchall(), [[7], [8], [9]])
+        self.assertEqual(rs.rownumber, 10)
+        self.assertEqual(offsets, [0, 3, 6, 9])
+
+        self.assertIsNone(rs.fetchone())
+        self.assertEqual(rs.rownumber, 10)
+        self.assertEqual(offsets, [0, 3, 6, 9])
+
+    def test_fetchone_refill_does_not_fetch_extra_pages(self):
+        # A refill must not make extra round-trips to fill ``arraysize`` when
+        # server pages are smaller (e.g. byte-limited): it serves what the current
+        # queue holds and only fetches a new page once the queue is empty.
+        batch_list = [[[0], [1]], [[2], [3]], [[4]]]
+        offsets = []
+        rs = self.make_dummy_result_set_from_batch_list(batch_list, offsets=offsets)
+        rs.arraysize = 100
+        self.assertEqual(offsets, [0])
+
+        self.assertSequenceEqual(rs.fetchone(), [0])
+        self.assertEqual(offsets, [0])
+        self.assertSequenceEqual(rs.fetchone(), [1])
+        self.assertEqual(offsets, [0])
+        self.assertSequenceEqual(rs.fetchone(), [2])
+        self.assertEqual(offsets, [0, 2])
+        self.assertEqual(rs.rownumber, 3)
+        self.assertEqualRowValues(rs.fetchall(), [[3], [4]])
+        self.assertEqual(offsets, [0, 2, 4])
+        self.assertIsNone(rs.fetchone())
+        self.assertEqual(rs.rownumber, 5)
+
     # Regression tests for fetchmany_arrow / fetchall_arrow handling of
     # the schemaless CloudFetch placeholder.
     def test_fetchall_arrow_drops_mismatched_empty_placeholder(self):
diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py
index cfa2d3f2e..66259c679 100644
--- a/tests/unit/test_kernel_result_set.py
+++ b/tests/unit/test_kernel_result_set.py
@@ -43,7 +43,7 @@ def close(self):
         self.closed = True
 
 
-def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet:
+def _make_rs(handle, *, disable_pandas=False, arraysize=100) -> KernelResultSet:
     # The base ResultSet __init__ takes a `connection` ref it never
     # actually dereferences during these buffer tests, so a Mock is
     # fine.
@@ -55,7 +55,7 @@ def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet:
         backend=backend,
         kernel_handle=handle,
         command_id=CommandId.from_sea_statement_id("smoke-test"),
-        arraysize=100,
+        arraysize=arraysize,
         buffer_size_bytes=1024,
     )
 
@@ -184,6 +184,69 @@ def test_fetchall_rows(int_schema):
     assert [r[0] for r in rows] == [1, 2, 3]
 
 
+def _ten_rows_handle(int_schema):
+    # Batch boundaries (3/4/3) deliberately don't line up with arraysize=4
+    # so fetchone() refills span kernel batches.
+    return _FakeKernelHandle(
+        int_schema,
+        [
+            _batch(int_schema, [0, 1, 2]),
+            _batch(int_schema, [3, 4, 5, 6]),
+            _batch(int_schema, [7, 8, 9]),
+        ],
+    )
+
+
+def test_mixed_fetches_keep_order_and_rownumber(int_schema):
+    rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
+    assert rs.fetchone()[0] == 0
+    assert rs.fetchone()[0] == 1
+    assert rs.rownumber == 2
+    assert [r[0] for r in rs.fetchmany(3)] == [2, 3, 4]
+    assert rs.rownumber == 5
+    assert rs.fetchone()[0] == 5
+    assert rs.rownumber == 6
+    assert [r[0] for r in rs.fetchall()] == [6, 7, 8, 9]
+    assert rs.rownumber == 10
+    assert rs.fetchone() is None
+    assert rs.rownumber == 10
+
+
+def test_fetchone_then_fetchall_arrow_keeps_order(int_schema):
+    rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
+    assert rs.fetchone()[0] == 0
+    assert rs.rownumber == 1
+    table = rs.fetchall_arrow()
+    assert table.column(0).to_pylist() == list(range(1, 10))
+    assert rs.rownumber == 10
+
+
+def test_fetchone_then_fetchmany_arrow_keeps_order(int_schema):
+    rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
+    assert rs.fetchone()[0] == 0
+    # Served entirely from the fetchone() row buffer (3 rows left in it).
+    assert rs.fetchmany_arrow(2).column(0).to_pylist() == [1, 2]
+    assert rs.rownumber == 3
+    # Drains the last buffered row, then continues from the batch FIFO.
+    assert rs.fetchmany_arrow(3).column(0).to_pylist() == [3, 4, 5]
+    assert rs.rownumber == 6
+    assert rs.fetchmany_arrow(0).num_rows == 0
+    assert rs.rownumber == 6
+    assert rs.fetchmany_arrow(10).column(0).to_pylist() == [6, 7, 8, 9]
+    assert rs.rownumber == 10
+
+
+def test_fetchone_then_fetchmany_rows_spans_buffer(int_schema):
+    rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
+    assert rs.fetchone()[0] == 0
+    # 3 rows come from the row buffer, the rest from the batch FIFO.
+    assert [r[0] for r in rs.fetchmany(6)] == [1, 2, 3, 4, 5, 6]
+    assert rs.rownumber == 7
+    assert [r[0] for r in rs.fetchmany(10)] == [7, 8, 9]
+    assert rs.rownumber == 10
+    assert rs.fetchmany(1) == []
+
+
 def test_fetchmany_negative_raises(int_schema):
     rs = _make_rs(_FakeKernelHandle(int_schema, []))
     with pytest.raises(ValueError):

From bf198d28e80ee1c207da54305fbbd78c126bb44a Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:12:47 +0000
Subject: [PATCH 04/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4418426389 at tests/unit/test_fetches.py:458
  - #4418426396 at tests/unit/test_sea_result_set.py:530

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 tests/unit/test_fetches.py        | 1 +
 tests/unit/test_sea_result_set.py | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 8726c59ac..ce97c7df9 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -534,6 +534,7 @@ def test_mixed_fetches_keep_order_and_rownumber(self):
         self.assertEqual(rs.rownumber, 10)
         self.assertIsNone(rs.fetchone())
 
+    @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
     def test_fetchone_then_fetchall_arrow_keeps_order(self):
         rs = self.make_columnar_result_set(10, arraysize=4)
         self.assertEqual(rs.fetchone()[0], 0)
diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py
index b004085a5..79771dbdb 100644
--- a/tests/unit/test_sea_result_set.py
+++ b/tests/unit/test_sea_result_set.py
@@ -527,6 +527,7 @@ def test_json_fetchone_converts_in_batches(
         assert [r.col2 for r in rows] == list(range(100))
         assert m.call_count <= 100 // 10 + 2
 
+    @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
     def test_json_mixed_fetches_keep_order_and_rownumber(
         self, mock_connection, mock_sea_client, execute_response
     ):
@@ -547,6 +548,7 @@ def test_json_mixed_fetches_keep_order_and_rownumber(
         assert rs.rownumber == 10
         assert rs.fetchone() is None
 
+    @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
     def test_json_fetchone_then_fetchall_arrow_keeps_order(
         self, mock_connection, mock_sea_client, execute_response
     ):

From 1b77e5757ec62c3874527131de386bd4d3c40e79 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:16:03 +0000
Subject: [PATCH 05/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4418426396 at tests/unit/test_sea_result_set.py:531

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 tests/unit/test_sea_result_set.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py
index 79771dbdb..56f45154f 100644
--- a/tests/unit/test_sea_result_set.py
+++ b/tests/unit/test_sea_result_set.py
@@ -527,6 +527,24 @@ def test_json_fetchone_converts_in_batches(
         assert [r.col2 for r in rows] == list(range(100))
         assert m.call_count <= 100 // 10 + 2
 
+    def test_json_row_fetches_keep_order_and_rownumber(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Mixed fetchone/fetchmany/fetchall ordering without any Arrow path."""
+        rs = self._make_json_result_set(
+            mock_connection, mock_sea_client, execute_response, 10, arraysize=4
+        )
+        assert rs.fetchone().col2 == 0
+        assert rs.fetchone().col2 == 1
+        assert rs.rownumber == 2
+        assert [r.col2 for r in rs.fetchmany(3)] == [2, 3, 4]
+        assert rs.rownumber == 5
+        assert rs.fetchone().col2 == 5
+        assert rs.rownumber == 6
+        assert [r.col2 for r in rs.fetchall()] == [6, 7, 8, 9]
+        assert rs.rownumber == 10
+        assert rs.fetchone() is None
+
     @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
     def test_json_mixed_fetches_keep_order_and_rownumber(
         self, mock_connection, mock_sea_client, execute_response

From bde5cf89c9bed1185fd40894e86fdc498a36c52e Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:16:30 +0000
Subject: [PATCH 06/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4418469118 at src/databricks/sql/result_set.py:216

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index fd0042983..c493a435a 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -215,6 +215,12 @@ def _fetchone_buffered(
         """
         if self._row_buffer_count() == 0:
             raw = fetch_raw(max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1))
+            # If convert() raises, this batch is dropped on purpose: the queue has
+            # already moved past it, and _next_row_index (advanced by fetch_raw)
+            # still equals the queue position. Rolling the index back would make it
+            # lag the queue and break Thrift's expected_row_start_offset check on
+            # the next fetch, and retrying convert() on the same data would just
+            # fail again.
             rows = convert(raw)
             if not rows:
                 return None

From af2c1fc845dd08e185952a696eb8aa4c9965228f Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:30:17 +0000
Subject: [PATCH 07/29] ai: apply changes for #1007 (4 review threads)

Addresses:
  - #4418502676 at src/databricks/sql/backend/kernel/result_set.py:229
  - #4418502680 at src/databricks/sql/result_set.py:100
  - #4418502685 at src/databricks/sql/result_set.py:414
  - #4418502697 at src/databricks/sql/backend/sea/result_set.py:202

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 .../sql/backend/kernel/result_set.py          | 16 +++-
 src/databricks/sql/result_set.py              | 18 +++--
 tests/unit/test_kernel_result_set.py          | 30 ++++++-
 tests/unit/test_sea_result_set.py             | 79 +++++++++++++++++++
 4 files changed, 131 insertions(+), 12 deletions(-)

diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index 800b9deea..d269d6454 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -197,6 +197,17 @@ def _fetchmany_arrow_unbuffered(self, size: int) -> pyarrow.Table:
         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 _prepend_buffered(
         buffered: Optional[pyarrow.Table], rest: pyarrow.Table
@@ -228,7 +239,7 @@ def fetchone(self) -> Optional[Row]:
         # 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._fetchmany_arrow_unbuffered, self._convert_arrow_table
+            self._fetch_available, self._convert_arrow_table
         )
 
     def fetchmany(self, size: int) -> List[Row]:
@@ -258,6 +269,9 @@ def close(self) -> None:
         connection-close would otherwise issue a kernel call into an
         already-disposed session.
         """
+        # Drop rows parked by fetchone() so post-close fetches return
+        # nothing, matching the cleared batch FIFO below.
+        self._row_buffer_clear()
         if self._kernel_handle is None:
             return
         if not self.connection.open:
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index c493a435a..221e61f8d 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -203,9 +203,9 @@ def _fetchone_buffered(
         """Serve fetchone() from the row buffer, refilling in small batches.
 
         Each refill asks ``fetch_raw`` for up to ``arraysize`` rows, capped at
-        ``_FETCHONE_BUFFER_MAX_ROWS``. ``fetch_raw`` may return fewer (Thrift only
-        serves what its current queue holds, refilling it at most until one row is
-        available, so latency to the first row is unchanged).
+        ``_FETCHONE_BUFFER_MAX_ROWS``. ``fetch_raw`` may return fewer (Thrift and the
+        kernel backend only serve what is already buffered, refilling at most until
+        one row is available, so latency to the first row is unchanged).
 
         ``_next_row_index`` keeps counting rows handed to the caller: rows parked in
         the buffer are subtracted on refill and added back as they are taken. The
@@ -283,6 +283,8 @@ def close(self) -> None:
             if isinstance(e.args[1], CursorAlreadyClosedError):
                 logger.info("Operation was canceled by a prior request")
         finally:
+            # Drop rows parked by fetchone() so post-close fetches return nothing.
+            self._row_buffer_clear()
             self.has_been_closed_server_side = True
             self.status = CommandState.CLOSED
 
@@ -411,7 +413,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
         An empty sequence is returned when no more rows are available.
         """
         if size < 0:
-            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+            raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
         buffered = self._row_buffer_take_raw(size)
         remaining = size - (buffered.num_rows if buffered is not None else 0)
         rest = self._fetchmany_arrow_from_queue(remaining) if remaining > 0 else None
@@ -425,7 +427,7 @@ def fetchmany_columnar(self, size: int):
         An empty sequence is returned when no more rows are available.
         """
         if size < 0:
-            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+            raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
         buffered = self._row_buffer_take_raw(size)
         remaining = size - (buffered.num_rows if buffered is not None else 0)
         rest = self._fetchmany_columnar_from_queue(remaining) if remaining > 0 else None
@@ -436,7 +438,7 @@ def fetchmany_columnar(self, size: int):
     def _fetchmany_arrow_from_queue(self, size: int) -> "pyarrow.Table":
         """Fetch up to ``size`` rows from the results queue, bypassing the row buffer."""
         if size < 0:
-            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+            raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
 
         # Hold 0-row chunks aside instead of appending them to ``partial_result_chunks``.
         # CloudFetchQueue may return a placeholder empty table whose schema does not
@@ -473,7 +475,7 @@ def _fetchmany_arrow_from_queue(self, size: int) -> "pyarrow.Table":
     def _fetchmany_columnar_from_queue(self, size: int):
         """Fetch up to ``size`` rows from the results queue, bypassing the row buffer."""
         if size < 0:
-            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+            raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
 
         results = self.results.next_n_rows(size)
         n_remaining_rows = size - results.num_rows
@@ -605,7 +607,7 @@ def fetchmany(self, size: int) -> List[Row]:
         An empty sequence is returned when no more rows are available.
         """
         if size < 0:
-            raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
+            raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
         buffered = self._row_buffer_take_rows(size)
         remaining = size - len(buffered)
         if buffered and remaining == 0:
diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py
index 66259c679..f9572b018 100644
--- a/tests/unit/test_kernel_result_set.py
+++ b/tests/unit/test_kernel_result_set.py
@@ -212,6 +212,20 @@ def test_mixed_fetches_keep_order_and_rownumber(int_schema):
     assert rs.rownumber == 10
 
 
+def test_fetchone_refill_pulls_at_most_one_batch(int_schema):
+    # First-row latency: a fetchone() refill must not loop fetch_next_batch()
+    # until arraysize rows are buffered.
+    handle = _ten_rows_handle(int_schema)
+    rs = _make_rs(handle, arraysize=100)
+    assert rs.fetchone()[0] == 0
+    assert len(handle._batches) == 2
+    assert rs.fetchone()[0] == 1
+    assert rs.fetchone()[0] == 2
+    assert len(handle._batches) == 2
+    assert rs.fetchone()[0] == 3
+    assert len(handle._batches) == 1
+
+
 def test_fetchone_then_fetchall_arrow_keeps_order(int_schema):
     rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
     assert rs.fetchone()[0] == 0
@@ -224,10 +238,11 @@ def test_fetchone_then_fetchall_arrow_keeps_order(int_schema):
 def test_fetchone_then_fetchmany_arrow_keeps_order(int_schema):
     rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
     assert rs.fetchone()[0] == 0
-    # Served entirely from the fetchone() row buffer (3 rows left in it).
+    # Served entirely from the fetchone() row buffer (2 rows left in it: the
+    # refill only took the first kernel batch).
     assert rs.fetchmany_arrow(2).column(0).to_pylist() == [1, 2]
     assert rs.rownumber == 3
-    # Drains the last buffered row, then continues from the batch FIFO.
+    # Row buffer is empty; continues from the batch FIFO.
     assert rs.fetchmany_arrow(3).column(0).to_pylist() == [3, 4, 5]
     assert rs.rownumber == 6
     assert rs.fetchmany_arrow(0).num_rows == 0
@@ -239,7 +254,7 @@ def test_fetchone_then_fetchmany_arrow_keeps_order(int_schema):
 def test_fetchone_then_fetchmany_rows_spans_buffer(int_schema):
     rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
     assert rs.fetchone()[0] == 0
-    # 3 rows come from the row buffer, the rest from the batch FIFO.
+    # 2 rows come from the row buffer, the rest from the batch FIFO.
     assert [r[0] for r in rs.fetchmany(6)] == [1, 2, 3, 4, 5, 6]
     assert rs.rownumber == 7
     assert [r[0] for r in rs.fetchmany(10)] == [7, 8, 9]
@@ -264,6 +279,15 @@ def test_close_is_idempotent_and_calls_handle(int_schema):
     rs.close()  # second call is a no-op (kernel handle is None)
 
 
+def test_close_drops_fetchone_row_buffer(int_schema):
+    # Rows parked in the fetchone() buffer must not survive close().
+    rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
+    assert rs.fetchone()[0] == 0
+    rs.close()
+    assert rs.fetchone() is None
+    assert rs.fetchall() == []
+
+
 def test_empty_stream(int_schema):
     rs = _make_rs(_FakeKernelHandle(int_schema, []))
     assert rs.fetchone() is None
diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py
index 56f45154f..6aecb1de8 100644
--- a/tests/unit/test_sea_result_set.py
+++ b/tests/unit/test_sea_result_set.py
@@ -579,6 +579,85 @@ def test_json_fetchone_then_fetchall_arrow_keeps_order(
         assert table.column(1).to_pylist() == list(range(1, 10))
         assert rs.rownumber == 10
 
+    def _make_arrow_result_set(
+        self, mock_connection, mock_sea_client, execute_response, n_rows, arraysize
+    ):
+        from databricks.sql.utils import ArrowQueue
+
+        mock_connection.disable_pandas = True
+        table = pyarrow.Table.from_pydict(
+            {
+                "col1": ["v{}".format(i) for i in range(n_rows)],
+                "col2": list(range(n_rows)),
+                "col3": [True] * n_rows,
+            }
+        )
+        result_data = ResultData(data=None, external_links=None, row_count=n_rows)
+        with patch(
+            "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue",
+            return_value=ArrowQueue(table, n_rows),
+        ):
+            return SeaResultSet(
+                connection=mock_connection,
+                execute_response=execute_response,
+                sea_client=mock_sea_client,
+                result_data=result_data,
+                manifest=self._create_empty_manifest(ResultFormat.ARROW_STREAM),
+                buffer_size_bytes=1000,
+                arraysize=arraysize,
+            )
+
+    @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
+    def test_arrow_mixed_fetches_keep_order_and_rownumber(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Buffered Arrow slice joined with the queue's next chunk via concat."""
+        rs = self._make_arrow_result_set(
+            mock_connection, mock_sea_client, execute_response, 10, arraysize=4
+        )
+        assert rs.fetchone().col2 == 0
+        assert rs.rownumber == 1
+        # 3 rows left in the buffer; ask for 5 so 2 must come from the queue.
+        table = rs.fetchmany_arrow(5)
+        assert isinstance(table, pyarrow.Table)
+        assert table.column(1).to_pylist() == [1, 2, 3, 4, 5]
+        assert rs.rownumber == 6
+        assert rs.fetchone().col2 == 6
+        assert rs.rownumber == 7
+        # Fully served from the buffer (remaining == 0 path).
+        assert rs.fetchmany_arrow(1).column(1).to_pylist() == [7]
+        assert rs.rownumber == 8
+        assert [r.col2 for r in rs.fetchall()] == [8, 9]
+        assert rs.rownumber == 10
+        assert rs.fetchone() is None
+
+    @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
+    def test_arrow_buffer_with_exhausted_queue_returns_buffer(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """Queue returns a 0-row slice after the buffer drains: buffered rows only."""
+        rs = self._make_arrow_result_set(
+            mock_connection, mock_sea_client, execute_response, 3, arraysize=10
+        )
+        assert rs.fetchone().col2 == 0
+        table = rs.fetchmany_arrow(5)
+        assert table.column(1).to_pylist() == [1, 2]
+        assert rs.rownumber == 3
+        assert rs.fetchall_arrow().num_rows == 0
+        assert rs.rownumber == 3
+
+    @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
+    def test_arrow_fetchone_then_fetchall_arrow_keeps_order(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        rs = self._make_arrow_result_set(
+            mock_connection, mock_sea_client, execute_response, 10, arraysize=4
+        )
+        assert rs.fetchone().col2 == 0
+        table = rs.fetchall_arrow()
+        assert table.column(1).to_pylist() == list(range(1, 10))
+        assert rs.rownumber == 10
+
     def test_is_staging_operation(
         self, mock_connection, mock_sea_client, execute_response
     ):

From 5f45f565efce14d65e57cfd37e691547e68cd5f3 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:34:33 +0000
Subject: [PATCH 08/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4418605049 at src/databricks/sql/result_set.py:224

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 20 ++++++++++++++++++--
 tests/unit/test_fetches.py       | 17 +++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 221e61f8d..e9d375a26 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -220,8 +220,24 @@ def _fetchone_buffered(
             # still equals the queue position. Rolling the index back would make it
             # lag the queue and break Thrift's expected_row_start_offset check on
             # the next fetch, and retrying convert() on the same data would just
-            # fail again.
-            rows = convert(raw)
+            # fail again. A warning records how many rows were discarded so a
+            # caller that swallows the error and keeps fetching isn't silently
+            # short of rows.
+            try:
+                rows = convert(raw)
+            except Exception:
+                if raw is None:
+                    dropped = 0
+                elif isinstance(raw, list):
+                    dropped = len(raw)
+                else:
+                    dropped = getattr(raw, "num_rows", 0)
+                logger.warning(
+                    "fetchone(): failed to convert a batch of %d row(s); these rows "
+                    "are discarded and will not be returned by later fetches",
+                    dropped,
+                )
+                raise
             if not rows:
                 return None
             self._row_buffer_raw = raw
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index ce97c7df9..0133459f4 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -534,6 +534,23 @@ def test_mixed_fetches_keep_order_and_rownumber(self):
         self.assertEqual(rs.rownumber, 10)
         self.assertIsNone(rs.fetchone())
 
+    def test_fetchone_convert_failure_logs_dropped_rows_and_reraises(self):
+        from unittest.mock import patch
+
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            side_effect=ValueError("bad value"),
+        ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
+            with self.assertRaises(ValueError):
+                rs.fetchone()
+        self.assertTrue(any("batch of 4 row(s)" in m for m in logs.output))
+        # The failed batch is dropped; the next fetch resumes after it.
+        self.assertEqual(rs._row_buffer_count(), 0)
+        self.assertEqual(rs.rownumber, 4)
+        self.assertEqual(rs.fetchone()[0], 4)
+
     @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
     def test_fetchone_then_fetchall_arrow_keeps_order(self):
         rs = self.make_columnar_result_set(10, arraysize=4)

From 1a44df9982b41150d097147a446cd61db2e792bb Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:44:07 +0000
Subject: [PATCH 09/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4418631883 at src/databricks/sql/backend/sea/result_set.py:250
  - #4418631888 at src/databricks/sql/result_set.py:222

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/backend/sea/result_set.py | 22 +++++-
 src/databricks/sql/result_set.py             | 77 ++++++++++++++------
 tests/unit/test_fetches.py                   | 38 ++++++++--
 tests/unit/test_sea_result_set.py            | 37 ++++++++++
 4 files changed, 146 insertions(+), 28 deletions(-)

diff --git a/src/databricks/sql/backend/sea/result_set.py b/src/databricks/sql/backend/sea/result_set.py
index babcd9c8b..f62f83f43 100644
--- a/src/databricks/sql/backend/sea/result_set.py
+++ b/src/databricks/sql/backend/sea/result_set.py
@@ -19,7 +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
+from databricks.sql.utils import CloudFetchQueue, concat_table_chunks
 
 logger = logging.getLogger(__name__)
 
@@ -186,6 +186,22 @@ def _fetch_raw_from_queue(self, size: Optional[int]) -> Any:
         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, ``next_n_rows`` keeps walking downloaded links until it has
+        ``size`` rows, so a refill that starts near the end of the current chunk
+        would block on the next download before returning a row that is already
+        available. Cap the request at what the current chunk still holds (matching
+        Thrift's ``_fetch_available_from_queue`` and the kernel's ``_fetch_available``).
+        Only called while the fetchone() row buffer is empty.
+        """
+        if isinstance(self.results, CloudFetchQueue) and self.results.table:
+            available = self.results.table.num_rows - self.results.table_row_index
+            if available > 0:
+                size = min(size, available)
+        return self._fetch_raw_from_queue(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)
@@ -247,7 +263,9 @@ def fetchone(self) -> Optional[Row]:
 
         if isinstance(self.results, JsonQueue):
             return self._fetchone_buffered(self.fetchmany_json, self._create_json_table)
-        return self._fetchone_buffered(self.fetchmany_arrow, self._convert_arrow_table)
+        return self._fetchone_buffered(
+            self._fetch_available_from_queue, self._convert_arrow_table
+        )
 
     def fetchmany(self, size: int) -> List[Row]:
         """
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index e9d375a26..5cbdb9821 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -34,6 +34,15 @@
 _FETCHONE_BUFFER_MAX_ROWS = 1000
 
 
+class _FailedRow:
+    """Placeholder in the fetchone() row buffer for a row that failed to convert."""
+
+    __slots__ = ("error",)
+
+    def __init__(self, error: Exception):
+        self.error = error
+
+
 class ResultSet(ABC):
     """
     Abstract base class for result sets returned by different backend implementations.
@@ -179,7 +188,44 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
         start, length = self._row_buffer_advance(n)
         if self._row_buffer_count() == 0:
             self._row_buffer_clear()
-        return rows[start : start + length]
+        taken = rows[start : start + length]
+        for row in taken:
+            if isinstance(row, _FailedRow):
+                # The taken rows are consumed either way, matching how a failed
+                # conversion behaved before fetchone() buffering.
+                raise row.error
+        return taken
+
+    @staticmethod
+    def _convert_rows_individually(
+        raw: Any, convert: Callable[[Any], List[Row]]
+    ) -> Optional[List[Any]]:
+        """Convert ``raw`` row by row after a batch conversion failed.
+
+        Rows that fail are returned as ``_FailedRow`` placeholders so the error is
+        raised only when that row is handed out. Returns None if ``raw`` has no rows
+        to split, in which case the caller re-raises the batch error.
+        """
+        if raw is None:
+            return None
+        n = len(raw) if isinstance(raw, list) else getattr(raw, "num_rows", 0)
+        if not n:
+            return None
+        rows: List[Any] = []
+        for i in range(n):
+            one = raw[i : i + 1] if isinstance(raw, list) else raw.slice(i, 1)
+            try:
+                rows.extend(convert(one))
+            except Exception as e:
+                rows.append(_FailedRow(e))
+        failed = sum(isinstance(r, _FailedRow) for r in rows)
+        logger.warning(
+            "fetchone(): failed to convert %d of %d row(s) in a batch; each failing "
+            "row raises when fetched",
+            failed,
+            n,
+        )
+        return rows
 
     def _row_buffer_take_raw(self, n: Optional[int] = None) -> Any:
         """Pop up to ``n`` (all if None) un-converted rows from the buffer.
@@ -215,29 +261,18 @@ def _fetchone_buffered(
         """
         if self._row_buffer_count() == 0:
             raw = fetch_raw(max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1))
-            # If convert() raises, this batch is dropped on purpose: the queue has
-            # already moved past it, and _next_row_index (advanced by fetch_raw)
-            # still equals the queue position. Rolling the index back would make it
-            # lag the queue and break Thrift's expected_row_start_offset check on
-            # the next fetch, and retrying convert() on the same data would just
-            # fail again. A warning records how many rows were discarded so a
-            # caller that swallows the error and keeps fetching isn't silently
-            # short of rows.
             try:
                 rows = convert(raw)
             except Exception:
-                if raw is None:
-                    dropped = 0
-                elif isinstance(raw, list):
-                    dropped = len(raw)
-                else:
-                    dropped = getattr(raw, "num_rows", 0)
-                logger.warning(
-                    "fetchone(): failed to convert a batch of %d row(s); these rows "
-                    "are discarded and will not be returned by later fetches",
-                    dropped,
-                )
-                raise
+                # The queue has already moved past this batch, so it can't be
+                # re-read. Fall back to converting it one row at a time so that,
+                # as before batching, only the row(s) that fail to convert are
+                # lost: each failing row is parked as a _FailedRow and re-raises
+                # its own error when it is reached.
+                fallback = self._convert_rows_individually(raw, convert)
+                if fallback is None:
+                    raise
+                rows = fallback
             if not rows:
                 return None
             self._row_buffer_raw = raw
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 0133459f4..73ae24d26 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -534,21 +534,49 @@ def test_mixed_fetches_keep_order_and_rownumber(self):
         self.assertEqual(rs.rownumber, 10)
         self.assertIsNone(rs.fetchone())
 
-    def test_fetchone_convert_failure_logs_dropped_rows_and_reraises(self):
+    def test_fetchone_convert_failure_loses_only_the_bad_row(self):
         from unittest.mock import patch
 
+        real_convert = ThriftResultSet._convert_columnar_table
+
+        def convert_failing_on_2(self_, table):
+            rows = real_convert(self_, table)
+            if any(r[0] == 2 for r in rows):
+                raise ValueError("bad value")
+            return rows
+
         rs = self.make_columnar_result_set(10, arraysize=4)
         with patch.object(
             ThriftResultSet,
             "_convert_columnar_table",
-            side_effect=ValueError("bad value"),
+            autospec=True,
+            side_effect=convert_failing_on_2,
         ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
+            self.assertEqual(rs.fetchone()[0], 0)
+            self.assertEqual(rs.fetchone()[0], 1)
             with self.assertRaises(ValueError):
                 rs.fetchone()
-        self.assertTrue(any("batch of 4 row(s)" in m for m in logs.output))
-        # The failed batch is dropped; the next fetch resumes after it.
+            # Only the failing row is lost; the rest of the batch is still served.
+            self.assertEqual(rs.rownumber, 3)
+            self.assertEqual(rs.fetchone()[0], 3)
+            self.assertEqual([r[0] for r in rs.fetchall()], list(range(4, 10)))
+        self.assertTrue(any("1 of 4 row(s)" in m for m in logs.output))
+        self.assertEqual(rs.rownumber, 10)
+
+    def test_fetchone_convert_failure_on_every_row_reraises_per_row(self):
+        from unittest.mock import patch
+
+        rs = self.make_columnar_result_set(6, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            side_effect=ValueError("bad value"),
+        ):
+            for expected_rownumber in range(1, 5):
+                with self.assertRaises(ValueError):
+                    rs.fetchone()
+                self.assertEqual(rs.rownumber, expected_rownumber)
         self.assertEqual(rs._row_buffer_count(), 0)
-        self.assertEqual(rs.rownumber, 4)
         self.assertEqual(rs.fetchone()[0], 4)
 
     @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py
index 6aecb1de8..31ebed864 100644
--- a/tests/unit/test_sea_result_set.py
+++ b/tests/unit/test_sea_result_set.py
@@ -658,6 +658,43 @@ def test_arrow_fetchone_then_fetchall_arrow_keeps_order(
         assert table.column(1).to_pylist() == list(range(1, 10))
         assert rs.rownumber == 10
 
+    @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
+    def test_cloudfetch_fetchone_refill_capped_to_current_chunk(
+        self, mock_connection, mock_sea_client, execute_response
+    ):
+        """fetchone() must not block on the next link while the current chunk has rows."""
+        from databricks.sql.backend.sea.queue import SeaCloudFetchQueue
+
+        chunk = pyarrow.Table.from_pydict(
+            {
+                "col1": ["v{}".format(i) for i in range(10)],
+                "col2": list(range(10)),
+                "col3": [True] * 10,
+            }
+        )
+        queue = Mock(spec=SeaCloudFetchQueue)
+        queue.table = chunk
+        queue.table_row_index = 8
+        queue.next_n_rows.side_effect = lambda n: chunk.slice(8, n)
+        result_data = ResultData(data=None, external_links=[], row_count=10)
+        with patch(
+            "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue",
+            return_value=queue,
+        ):
+            rs = SeaResultSet(
+                connection=mock_connection,
+                execute_response=execute_response,
+                sea_client=mock_sea_client,
+                result_data=result_data,
+                manifest=self._create_empty_manifest(ResultFormat.ARROW_STREAM),
+                buffer_size_bytes=1000,
+                arraysize=100,
+            )
+        assert rs.fetchone().col2 == 8
+        queue.next_n_rows.assert_called_once_with(2)
+        assert rs.fetchone().col2 == 9
+        assert rs.rownumber == 2
+
     def test_is_staging_operation(
         self, mock_connection, mock_sea_client, execute_response
     ):

From 05f9f568ecbf6215b4cc80aec2108acb588e1499 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 22:55:30 +0000
Subject: [PATCH 10/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4418695007 at src/databricks/sql/result_set.py:557
  - #4418695015 at src/databricks/sql/result_set.py:281

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/backend/sea/result_set.py |  2 +-
 src/databricks/sql/result_set.py             | 41 +++++++++++++++++---
 tests/unit/test_fetches.py                   | 40 +++++++++++++++++++
 3 files changed, 76 insertions(+), 7 deletions(-)

diff --git a/src/databricks/sql/backend/sea/result_set.py b/src/databricks/sql/backend/sea/result_set.py
index f62f83f43..110cd9cc9 100644
--- a/src/databricks/sql/backend/sea/result_set.py
+++ b/src/databricks/sql/backend/sea/result_set.py
@@ -193,7 +193,7 @@ def _fetch_available_from_queue(self, size: int) -> Any:
         ``size`` rows, so a refill that starts near the end of the current chunk
         would block on the next download before returning a row that is already
         available. Cap the request at what the current chunk still holds (matching
-        Thrift's ``_fetch_available_from_queue`` and the kernel's ``_fetch_available``).
+        Thrift's ``_cap_to_current_chunk`` and the kernel's ``_fetch_available``).
         Only called while the fetchone() row buffer is empty.
         """
         if isinstance(self.results, CloudFetchQueue) and self.results.table:
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 5cbdb9821..8a6c756aa 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -16,8 +16,9 @@
     from databricks.sql.client import Connection
 from databricks.sql.backend.databricks_client import DatabricksClient
 from databricks.sql.types import Row
-from databricks.sql.exc import RequestError, CursorAlreadyClosedError
+from databricks.sql.exc import InternalError, RequestError, CursorAlreadyClosedError
 from databricks.sql.utils import (
+    CloudFetchQueue,
     ColumnTable,
     ColumnQueue,
     concat_table_chunks,
@@ -196,6 +197,14 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
                 raise row.error
         return taken
 
+    @staticmethod
+    def _raw_row_count(raw: Any) -> Optional[int]:
+        """Row count of a raw chunk (Arrow table, ColumnTable or list); None if unknown."""
+        if raw is None:
+            return 0
+        n = len(raw) if isinstance(raw, list) else getattr(raw, "num_rows", None)
+        return n if isinstance(n, int) else None
+
     @staticmethod
     def _convert_rows_individually(
         raw: Any, convert: Callable[[Any], List[Row]]
@@ -206,9 +215,7 @@ def _convert_rows_individually(
         raised only when that row is handed out. Returns None if ``raw`` has no rows
         to split, in which case the caller re-raises the batch error.
         """
-        if raw is None:
-            return None
-        n = len(raw) if isinstance(raw, list) else getattr(raw, "num_rows", 0)
+        n = ResultSet._raw_row_count(raw)
         if not n:
             return None
         rows: List[Any] = []
@@ -275,6 +282,16 @@ def _fetchone_buffered(
                 rows = fallback
             if not rows:
                 return None
+            raw_count = self._raw_row_count(raw)
+            if raw_count is not None and raw_count != len(rows):
+                # fetch_raw advanced _next_row_index (and the Thrift fetch offset)
+                # by raw_count; the buffer must line up 1:1 with the raw chunk or
+                # rownumber and fetchmany_arrow slicing would silently drift.
+                raise InternalError(
+                    "fetchone(): converted {} row(s) from a chunk of {} raw row(s)".format(
+                        len(rows), raw_count
+                    )
+                )
             self._row_buffer_raw = raw
             self._row_buffer_rows = rows
             self._row_buffer_pos = 0
@@ -552,19 +569,31 @@ def _fetch_available_from_queue(self, size: int):
         until ``size`` rows are collected: it only refills the queue while it is
         empty, so fetchone() makes no more round-trips than it needs to return one
         row (server pages can be smaller than ``size`` when they hit the byte limit).
+        On CloudFetch, ``next_n_rows`` keeps walking downloaded files until it has
+        ``size`` rows, so the request is also capped at what the current chunk still
+        holds; otherwise a refill near a chunk boundary would block on the next
+        download before returning a row that is already available.
         Only called while the fetchone() row buffer is empty.
         """
-        results = self.results.next_n_rows(size)
+        results = self.results.next_n_rows(self._cap_to_current_chunk(size))
         while (
             results.num_rows == 0
             and not self.has_been_closed_server_side
             and self.has_more_rows
         ):
             self._fill_results_buffer()
-            results = self.results.next_n_rows(size)
+            results = self.results.next_n_rows(self._cap_to_current_chunk(size))
         self._next_row_index += results.num_rows
         return results
 
+    def _cap_to_current_chunk(self, size: int) -> int:
+        """Cap ``size`` at the rows left in the current CloudFetch chunk, if any."""
+        if isinstance(self.results, CloudFetchQueue) and self.results.table:
+            available = self.results.table.num_rows - self.results.table_row_index
+            if available > 0:
+                return min(size, available)
+        return size
+
     def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
         buffered = self._row_buffer_take_raw()
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 73ae24d26..f7e0b4c1e 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -392,6 +392,27 @@ def test_fetchone_refill_does_not_fetch_extra_pages(self):
         self.assertIsNone(rs.fetchone())
         self.assertEqual(rs.rownumber, 5)
 
+    def test_cloudfetch_fetchone_refill_capped_to_current_chunk(self):
+        # fetchone() must not block on the next CloudFetch file while the current
+        # chunk still holds rows: the refill is capped at what the chunk has left.
+        from databricks.sql.utils import ThriftCloudFetchQueue
+
+        _, chunk = self.make_arrow_table([[i] for i in range(10)])
+        queue = Mock(spec=ThriftCloudFetchQueue)
+        queue.table = chunk
+        queue.table_row_index = 8
+        queue.next_n_rows.side_effect = lambda n: chunk.slice(8, n)
+        rs = self.make_dummy_result_set_from_queue_list(
+            [queue],
+            description=[("col0", "integer", None, None, None, None, None)],
+        )
+        rs.arraysize = 100
+
+        self.assertSequenceEqual(rs.fetchone(), [8])
+        queue.next_n_rows.assert_called_once_with(2)
+        self.assertSequenceEqual(rs.fetchone(), [9])
+        self.assertEqual(rs.rownumber, 2)
+
     # Regression tests for fetchmany_arrow / fetchall_arrow handling of
     # the schemaless CloudFetch placeholder.
     def test_fetchall_arrow_drops_mismatched_empty_placeholder(self):
@@ -579,6 +600,25 @@ def test_fetchone_convert_failure_on_every_row_reraises_per_row(self):
         self.assertEqual(rs._row_buffer_count(), 0)
         self.assertEqual(rs.fetchone()[0], 4)
 
+    def test_fetchone_raw_converted_count_mismatch_fails_loudly(self):
+        from unittest.mock import patch
+        from databricks.sql.exc import InternalError
+
+        real_convert = ThriftResultSet._convert_columnar_table
+
+        def convert_dropping_a_row(self_, table):
+            return real_convert(self_, table)[:-1]
+
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert_dropping_a_row,
+        ):
+            with self.assertRaises(InternalError):
+                rs.fetchone()
+
     @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
     def test_fetchone_then_fetchall_arrow_keeps_order(self):
         rs = self.make_columnar_result_set(10, arraysize=4)

From 48a432bc1cf52fab3ca2fbe4ac9688f0cfbf0166 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 23:04:28 +0000
Subject: [PATCH 11/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4418777541 at src/databricks/sql/result_set.py:276

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 32 ++++++++++++++++++++++++++------
 tests/unit/test_fetches.py       | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 8a6c756aa..517be7b43 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -34,6 +34,12 @@
 # return one row.
 _FETCHONE_BUFFER_MAX_ROWS = 1000
 
+# After a fetchone() batch conversion fails, the batch is retried one row at a time.
+# If this many leading rows all fail too, the failure is treated as systematic: the
+# rest of the batch is not retried row by row (each would cost a separate
+# conversion) and those rows re-raise the original batch error instead.
+_FETCHONE_FALLBACK_PROBE_ROWS = 5
+
 
 class _FailedRow:
     """Placeholder in the fetchone() row buffer for a row that failed to convert."""
@@ -207,19 +213,29 @@ def _raw_row_count(raw: Any) -> Optional[int]:
 
     @staticmethod
     def _convert_rows_individually(
-        raw: Any, convert: Callable[[Any], List[Row]]
+        raw: Any, convert: Callable[[Any], List[Row]], batch_error: Exception
     ) -> Optional[List[Any]]:
         """Convert ``raw`` row by row after a batch conversion failed.
 
         Rows that fail are returned as ``_FailedRow`` placeholders so the error is
-        raised only when that row is handed out. Returns None if ``raw`` has no rows
-        to split, in which case the caller re-raises the batch error.
+        raised only when that row is handed out. If the first
+        ``_FETCHONE_FALLBACK_PROBE_ROWS`` rows all fail, the failure is treated as
+        systematic and the remaining rows are parked with ``batch_error`` instead of
+        being converted one by one. Returns None if ``raw`` has no rows to split, in
+        which case the caller re-raises the batch error.
         """
         n = ResultSet._raw_row_count(raw)
         if not n:
             return None
         rows: List[Any] = []
+        systematic = False
         for i in range(n):
+            if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(
+                isinstance(r, _FailedRow) for r in rows
+            ):
+                systematic = True
+                rows.extend(_FailedRow(batch_error) for _ in range(n - i))
+                break
             one = raw[i : i + 1] if isinstance(raw, list) else raw.slice(i, 1)
             try:
                 rows.extend(convert(one))
@@ -227,10 +243,14 @@ def _convert_rows_individually(
                 rows.append(_FailedRow(e))
         failed = sum(isinstance(r, _FailedRow) for r in rows)
         logger.warning(
-            "fetchone(): failed to convert %d of %d row(s) in a batch; each failing "
+            "fetchone(): failed to convert %d of %d row(s) in a batch%s; each failing "
             "row raises when fetched",
             failed,
             n,
+            " (first %d rows all failed, rest not retried)" % _FETCHONE_FALLBACK_PROBE_ROWS
+            if systematic
+            else "",
+            exc_info=batch_error,
         )
         return rows
 
@@ -270,13 +290,13 @@ def _fetchone_buffered(
             raw = fetch_raw(max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1))
             try:
                 rows = convert(raw)
-            except Exception:
+            except Exception as batch_error:
                 # The queue has already moved past this batch, so it can't be
                 # re-read. Fall back to converting it one row at a time so that,
                 # as before batching, only the row(s) that fail to convert are
                 # lost: each failing row is parked as a _FailedRow and re-raises
                 # its own error when it is reached.
-                fallback = self._convert_rows_individually(raw, convert)
+                fallback = self._convert_rows_individually(raw, convert, batch_error)
                 if fallback is None:
                     raise
                 rows = fallback
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index f7e0b4c1e..8f704dec2 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -600,6 +600,38 @@ def test_fetchone_convert_failure_on_every_row_reraises_per_row(self):
         self.assertEqual(rs._row_buffer_count(), 0)
         self.assertEqual(rs.fetchone()[0], 4)
 
+    def test_fetchone_systematic_convert_failure_stops_per_row_retries(self):
+        from unittest.mock import patch
+        from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS
+
+        batch_error = ValueError("batch failure")
+        calls = []
+
+        def convert(self_, table):
+            calls.append(table.num_rows)
+            if len(calls) == 1:
+                raise batch_error
+            raise ValueError("row failure")
+
+        rs = self.make_columnar_result_set(50, arraysize=20)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert,
+        ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
+            for i in range(20):
+                with self.assertRaises(ValueError) as ctx:
+                    rs.fetchone()
+                if i >= _FETCHONE_FALLBACK_PROBE_ROWS:
+                    self.assertIs(ctx.exception, batch_error)
+                self.assertEqual(rs.rownumber, i + 1)
+        # One batch attempt plus only the probe rows were converted individually.
+        self.assertEqual(len(calls), 1 + _FETCHONE_FALLBACK_PROBE_ROWS)
+        self.assertTrue(any("20 of 20 row(s)" in m for m in logs.output))
+        self.assertTrue(any("batch failure" in m for m in logs.output))
+        self.assertEqual(rs.fetchone()[0], 20)
+
     def test_fetchone_raw_converted_count_mismatch_fails_loudly(self):
         from unittest.mock import patch
         from databricks.sql.exc import InternalError

From 4b590fa703a13acf09a5498b1a19213f80538711 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 23:16:46 +0000
Subject: [PATCH 12/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4418836698 at src/databricks/sql/result_set.py:303
  - #4418836707 at src/databricks/sql/backend/sea/result_set.py:189

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/backend/sea/result_set.py | 11 +++-----
 src/databricks/sql/result_set.py             | 27 ++++++++++++--------
 tests/unit/test_fetches.py                   | 14 ++++++++++
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/src/databricks/sql/backend/sea/result_set.py b/src/databricks/sql/backend/sea/result_set.py
index 110cd9cc9..eb9b4b853 100644
--- a/src/databricks/sql/backend/sea/result_set.py
+++ b/src/databricks/sql/backend/sea/result_set.py
@@ -19,7 +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 CloudFetchQueue, concat_table_chunks
+from databricks.sql.utils import concat_table_chunks
 
 logger = logging.getLogger(__name__)
 
@@ -193,14 +193,11 @@ def _fetch_available_from_queue(self, size: int) -> Any:
         ``size`` rows, so a refill that starts near the end of the current chunk
         would block on the next download before returning a row that is already
         available. Cap the request at what the current chunk still holds (matching
-        Thrift's ``_cap_to_current_chunk`` and the kernel's ``_fetch_available``).
+        Thrift and the kernel's ``_fetch_available``) via the shared
+        ``ResultSet._cap_to_current_chunk``.
         Only called while the fetchone() row buffer is empty.
         """
-        if isinstance(self.results, CloudFetchQueue) and self.results.table:
-            available = self.results.table.num_rows - self.results.table_row_index
-            if available > 0:
-                size = min(size, available)
-        return self._fetch_raw_from_queue(size)
+        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."""
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 517be7b43..80ad2774c 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -254,6 +254,18 @@ def _convert_rows_individually(
         )
         return rows
 
+    def _cap_to_current_chunk(self, size: int) -> int:
+        """Cap ``size`` at the rows left in the current CloudFetch chunk, if any.
+
+        Shared by the Thrift and SEA fetchone() refills so the CloudFetchQueue
+        internals (``table`` / ``table_row_index``) are read in one place.
+        """
+        if isinstance(self.results, CloudFetchQueue) and self.results.table:
+            available = self.results.table.num_rows - self.results.table_row_index
+            if available > 0:
+                return min(size, available)
+        return size
+
     def _row_buffer_take_raw(self, n: Optional[int] = None) -> Any:
         """Pop up to ``n`` (all if None) un-converted rows from the buffer.
 
@@ -300,8 +312,9 @@ def _fetchone_buffered(
                 if fallback is None:
                     raise
                 rows = fallback
-            if not rows:
-                return None
+            # Check for drift before the empty-result early return, so a converter
+            # that yields [] for a non-empty chunk raises instead of silently
+            # ending iteration and losing those rows.
             raw_count = self._raw_row_count(raw)
             if raw_count is not None and raw_count != len(rows):
                 # fetch_raw advanced _next_row_index (and the Thrift fetch offset)
@@ -312,6 +325,8 @@ def _fetchone_buffered(
                         len(rows), raw_count
                     )
                 )
+            if not rows:
+                return None
             self._row_buffer_raw = raw
             self._row_buffer_rows = rows
             self._row_buffer_pos = 0
@@ -606,14 +621,6 @@ def _fetch_available_from_queue(self, size: int):
         self._next_row_index += results.num_rows
         return results
 
-    def _cap_to_current_chunk(self, size: int) -> int:
-        """Cap ``size`` at the rows left in the current CloudFetch chunk, if any."""
-        if isinstance(self.results, CloudFetchQueue) and self.results.table:
-            available = self.results.table.num_rows - self.results.table_row_index
-            if available > 0:
-                return min(size, available)
-        return size
-
     def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
         buffered = self._row_buffer_take_raw()
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 8f704dec2..9d9fa238e 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -651,6 +651,20 @@ def convert_dropping_a_row(self_, table):
             with self.assertRaises(InternalError):
                 rs.fetchone()
 
+    def test_fetchone_empty_conversion_of_nonempty_chunk_fails_loudly(self):
+        from unittest.mock import patch
+        from databricks.sql.exc import InternalError
+
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            return_value=[],
+        ):
+            with self.assertRaises(InternalError):
+                rs.fetchone()
+
     @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
     def test_fetchone_then_fetchall_arrow_keeps_order(self):
         rs = self.make_columnar_result_set(10, arraysize=4)

From 9e4320f04e883782599ee84307206a2c8344c8e9 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 23:21:48 +0000
Subject: [PATCH 13/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4418907858 at src/databricks/sql/result_set.py:613

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/backend/sea/result_set.py | 11 +++++------
 src/databricks/sql/result_set.py             | 12 ++++++++----
 tests/unit/test_fetches.py                   |  5 +++--
 tests/unit/test_sea_result_set.py            |  2 +-
 4 files changed, 17 insertions(+), 13 deletions(-)

diff --git a/src/databricks/sql/backend/sea/result_set.py b/src/databricks/sql/backend/sea/result_set.py
index eb9b4b853..09b53e842 100644
--- a/src/databricks/sql/backend/sea/result_set.py
+++ b/src/databricks/sql/backend/sea/result_set.py
@@ -189,12 +189,11 @@ def _fetch_raw_from_queue(self, size: Optional[int]) -> Any:
     def _fetch_available_from_queue(self, size: int) -> Any:
         """Fetch up to ``size`` rows the queue already holds, for Arrow fetchone() refills.
 
-        On CloudFetch, ``next_n_rows`` keeps walking downloaded links until it has
-        ``size`` rows, so a refill that starts near the end of the current chunk
-        would block on the next download before returning a row that is already
-        available. Cap the request at what the current chunk still holds (matching
-        Thrift and the kernel's ``_fetch_available``) via the shared
-        ``ResultSet._cap_to_current_chunk``.
+        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))
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 80ad2774c..989daf43b 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -257,6 +257,9 @@ def _convert_rows_individually(
     def _cap_to_current_chunk(self, size: int) -> int:
         """Cap ``size`` at the rows left in the current CloudFetch chunk, if any.
 
+        Limits each fetchone() refill to a single CloudFetch chunk. It does not
+        avoid waiting on the next download: ``next_n_rows`` advances to the next
+        file as soon as the current one is fully read, inside the same call.
         Shared by the Thrift and SEA fetchone() refills so the CloudFetchQueue
         internals (``table`` / ``table_row_index``) are read in one place.
         """
@@ -604,10 +607,11 @@ def _fetch_available_from_queue(self, size: int):
         until ``size`` rows are collected: it only refills the queue while it is
         empty, so fetchone() makes no more round-trips than it needs to return one
         row (server pages can be smaller than ``size`` when they hit the byte limit).
-        On CloudFetch, ``next_n_rows`` keeps walking downloaded files until it has
-        ``size`` rows, so the request is also capped at what the current chunk still
-        holds; otherwise a refill near a chunk boundary would block on the next
-        download before returning a row that is already available.
+        On CloudFetch, the request is also capped at what the current chunk still
+        holds (see ``_cap_to_current_chunk``), so a refill never stitches rows from
+        several downloaded files into one batch. This only bounds the batch: reading
+        to the end of a chunk still makes ``next_n_rows`` advance to (and wait for)
+        the next file before it returns, exactly as the uncapped path did.
         Only called while the fetchone() row buffer is empty.
         """
         results = self.results.next_n_rows(self._cap_to_current_chunk(size))
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 9d9fa238e..ab3d14399 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -393,8 +393,9 @@ def test_fetchone_refill_does_not_fetch_extra_pages(self):
         self.assertEqual(rs.rownumber, 5)
 
     def test_cloudfetch_fetchone_refill_capped_to_current_chunk(self):
-        # fetchone() must not block on the next CloudFetch file while the current
-        # chunk still holds rows: the refill is capped at what the chunk has left.
+        # The fetchone() refill is capped at the rows left in the current CloudFetch
+        # chunk, so one refill never spans several files. (This only bounds the
+        # batch size; the mock queue does not model advancing to the next file.)
         from databricks.sql.utils import ThriftCloudFetchQueue
 
         _, chunk = self.make_arrow_table([[i] for i in range(10)])
diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py
index 31ebed864..f78ee2ba9 100644
--- a/tests/unit/test_sea_result_set.py
+++ b/tests/unit/test_sea_result_set.py
@@ -662,7 +662,7 @@ def test_arrow_fetchone_then_fetchall_arrow_keeps_order(
     def test_cloudfetch_fetchone_refill_capped_to_current_chunk(
         self, mock_connection, mock_sea_client, execute_response
     ):
-        """fetchone() must not block on the next link while the current chunk has rows."""
+        """fetchone() refill is capped at the rows left in the current chunk (batch size only)."""
         from databricks.sql.backend.sea.queue import SeaCloudFetchQueue
 
         chunk = pyarrow.Table.from_pydict(

From 8e03ea6b1de2428f3e83a60cfefcbfd370d32edd Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 23:25:21 +0000
Subject: [PATCH 14/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4418928139 at tests/e2e/test_kernel_backend.py:209

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 tests/e2e/test_driver.py         | 5 ++++-
 tests/e2e/test_kernel_backend.py | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py
index ea99de595..47aa85979 100644
--- a/tests/e2e/test_driver.py
+++ b/tests/e2e/test_driver.py
@@ -992,7 +992,10 @@ def _assert_rows_and_bound(self, rows, calls, arraysize):
         assert len(rows) == n
         assert [r[0] for r in rows] == list(range(n))
         assert all(r[1] == str(r[0]) for r in rows)
-        bound = -(-n // arraysize) + 2
+        # Refills are capped per server batch / CloudFetch chunk, so an exact
+        # ceil(n/arraysize) bound is layout-dependent. Only distinguish "a few"
+        # conversions from "~n" (one per row, issue #981).
+        bound = n // 10
         assert calls <= bound, (
             "row conversion called {} times for {} rows (arraysize={}); "
             "expected at most {} (issue #981)".format(calls, n, arraysize, bound)
diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py
index 2b67f7df2..e1259af0e 100644
--- a/tests/e2e/test_kernel_backend.py
+++ b/tests/e2e/test_kernel_backend.py
@@ -206,7 +206,9 @@ def _assert_rows_and_bound(rows, calls):
     assert len(rows) == _ROW_BUFFER_N
     assert [r[0] for r in rows] == list(range(_ROW_BUFFER_N))
     assert all(r[1] == str(r[0]) for r in rows)
-    bound = -(-_ROW_BUFFER_N // _ROW_BUFFER_ARRAYSIZE) + 2
+    # Server batch boundaries add short refills, so an exact ceil(N/arraysize)
+    # bound is layout-dependent. Only distinguish "a few" from "~N" (issue #981).
+    bound = _ROW_BUFFER_N // 10
     assert calls <= bound, (
         f"_convert_arrow_table called {calls} times for {_ROW_BUFFER_N} rows; "
         f"expected at most {bound} (issue #981)"

From 5e2472608827e9890677158c64c2a11ea8cb8537 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 23:35:18 +0000
Subject: [PATCH 15/29] ai: apply changes for #1007 (3 review threads)

Addresses:
  - #4418954338 at src/databricks/sql/result_set.py:228
  - #4418954347 at src/databricks/sql/result_set.py:203
  - #4418954351 at src/databricks/sql/result_set.py:316

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 45 +++++++++++++++++++++++---------
 tests/unit/test_fetches.py       | 42 ++++++++++++++++++++++++++++-
 2 files changed, 73 insertions(+), 14 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 989daf43b..f462656b8 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -35,19 +35,25 @@
 _FETCHONE_BUFFER_MAX_ROWS = 1000
 
 # After a fetchone() batch conversion fails, the batch is retried one row at a time.
-# If this many leading rows all fail too, the failure is treated as systematic: the
-# rest of the batch is not retried row by row (each would cost a separate
-# conversion) and those rows re-raise the original batch error instead.
+# If this many leading rows all fail too, with the same error type and message as
+# the batch, the failure is treated as systematic: the rest of the batch is not
+# retried row by row (each would cost a separate conversion) and those rows re-raise
+# the original batch error instead. Leading rows failing for other reasons do not
+# trigger the shortcut, so later good rows are still converted.
 _FETCHONE_FALLBACK_PROBE_ROWS = 5
 
 
 class _FailedRow:
     """Placeholder in the fetchone() row buffer for a row that failed to convert."""
 
-    __slots__ = ("error",)
+    __slots__ = ("error", "traceback")
 
     def __init__(self, error: Exception):
         self.error = error
+        # Snapshot the conversion traceback. Systematic failures park the same
+        # exception instance in many rows; restoring this before each raise keeps
+        # its __traceback__ from growing (and pinning frames) on every fetchone().
+        self.traceback = error.__traceback__
 
 
 class ResultSet(ABC):
@@ -200,7 +206,7 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
             if isinstance(row, _FailedRow):
                 # The taken rows are consumed either way, matching how a failed
                 # conversion behaved before fetchone() buffering.
-                raise row.error
+                raise row.error.with_traceback(row.traceback)
         return taken
 
     @staticmethod
@@ -219,20 +225,29 @@ def _convert_rows_individually(
 
         Rows that fail are returned as ``_FailedRow`` placeholders so the error is
         raised only when that row is handed out. If the first
-        ``_FETCHONE_FALLBACK_PROBE_ROWS`` rows all fail, the failure is treated as
-        systematic and the remaining rows are parked with ``batch_error`` instead of
-        being converted one by one. Returns None if ``raw`` has no rows to split, in
+        ``_FETCHONE_FALLBACK_PROBE_ROWS`` rows all fail with the same error as the
+        batch (same type and message), the failure is treated as systematic and the
+        remaining rows are parked with ``batch_error`` instead of being converted one
+        by one. Otherwise every row is tried individually. Returns None if ``raw`` has no rows to split, in
         which case the caller re-raises the batch error.
         """
         n = ResultSet._raw_row_count(raw)
         if not n:
             return None
+        def same_as_batch(r: Any) -> bool:
+            return (
+                isinstance(r, _FailedRow)
+                and type(r.error) is type(batch_error)
+                and str(r.error) == str(batch_error)
+            )
+
         rows: List[Any] = []
         systematic = False
         for i in range(n):
-            if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(
-                isinstance(r, _FailedRow) for r in rows
-            ):
+            # Only skip the remaining per-row retries when every probe row failed
+            # with the same error as the batch; leading rows that fail for their
+            # own reasons must not cause later good rows to be dropped.
+            if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(same_as_batch(r) for r in rows):
                 systematic = True
                 rows.extend(_FailedRow(batch_error) for _ in range(n - i))
                 break
@@ -242,16 +257,20 @@ def _convert_rows_individually(
             except Exception as e:
                 rows.append(_FailedRow(e))
         failed = sum(isinstance(r, _FailedRow) for r in rows)
+        # One summary line, no traceback: the error itself still reaches the caller
+        # when the failing row is fetched. The full traceback is only logged at DEBUG.
         logger.warning(
             "fetchone(): failed to convert %d of %d row(s) in a batch%s; each failing "
-            "row raises when fetched",
+            "row raises when fetched (%s: %s)",
             failed,
             n,
             " (first %d rows all failed, rest not retried)" % _FETCHONE_FALLBACK_PROBE_ROWS
             if systematic
             else "",
-            exc_info=batch_error,
+            type(batch_error).__name__,
+            batch_error,
         )
+        logger.debug("fetchone(): batch conversion error", exc_info=batch_error)
         return rows
 
     def _cap_to_current_chunk(self, size: int) -> int:
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index ab3d14399..0d8b7e7bd 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -612,7 +612,8 @@ def convert(self_, table):
             calls.append(table.num_rows)
             if len(calls) == 1:
                 raise batch_error
-            raise ValueError("row failure")
+            # Same type and message as the batch error: a systematic failure.
+            raise ValueError("batch failure")
 
         rs = self.make_columnar_result_set(50, arraysize=20)
         with patch.object(
@@ -621,18 +622,57 @@ def convert(self_, table):
             autospec=True,
             side_effect=convert,
         ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
+            tb_depths = []
             for i in range(20):
                 with self.assertRaises(ValueError) as ctx:
                     rs.fetchone()
                 if i >= _FETCHONE_FALLBACK_PROBE_ROWS:
                     self.assertIs(ctx.exception, batch_error)
+                    tb, depth = ctx.exception.__traceback__, 0
+                    while tb is not None:
+                        tb, depth = tb.tb_next, depth + 1
+                    tb_depths.append(depth)
                 self.assertEqual(rs.rownumber, i + 1)
+            # Re-raising the shared batch error must not grow its traceback per row.
+            self.assertEqual(len(set(tb_depths)), 1)
         # One batch attempt plus only the probe rows were converted individually.
         self.assertEqual(len(calls), 1 + _FETCHONE_FALLBACK_PROBE_ROWS)
         self.assertTrue(any("20 of 20 row(s)" in m for m in logs.output))
         self.assertTrue(any("batch failure" in m for m in logs.output))
         self.assertEqual(rs.fetchone()[0], 20)
 
+    def test_fetchone_leading_bad_rows_do_not_drop_later_good_rows(self):
+        from unittest.mock import patch
+        from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS
+
+        real_convert = ThriftResultSet._convert_columnar_table
+        n_bad = _FETCHONE_FALLBACK_PROBE_ROWS + 1
+
+        def convert(self_, table):
+            rows = real_convert(self_, table)
+            bad = [r[0] for r in rows if r[0] < n_bad]
+            if bad:
+                # Each bad row fails with its own message, unlike the batch error.
+                raise ValueError("bad value in row {}".format(bad[0]))
+            return rows
+
+        rs = self.make_columnar_result_set(20, arraysize=20)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert,
+        ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
+            for i in range(n_bad):
+                with self.assertRaises(ValueError) as ctx:
+                    rs.fetchone()
+                self.assertEqual(str(ctx.exception), "bad value in row {}".format(i))
+            # Rows after the leading bad run still convert and are served.
+            self.assertEqual([r[0] for r in rs.fetchall()], list(range(n_bad, 20)))
+        self.assertTrue(any("{} of 20 row(s)".format(n_bad) in m for m in logs.output))
+        self.assertFalse(any("rest not retried" in m for m in logs.output))
+        self.assertEqual(rs.rownumber, 20)
+
     def test_fetchone_raw_converted_count_mismatch_fails_loudly(self):
         from unittest.mock import patch
         from databricks.sql.exc import InternalError

From 3cddc9a2323b530c101ed49e721074a0a651290c Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 23:45:09 +0000
Subject: [PATCH 16/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4419018676 at src/databricks/sql/result_set.py:231
  - #4419018680 at src/databricks/sql/result_set.py:265

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 49 +++++++++++++++++++++++++-------
 tests/unit/test_fetches.py       | 23 +++++++++++++++
 2 files changed, 61 insertions(+), 11 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index f462656b8..60aab5eb3 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -122,6 +122,8 @@ def __init__(
         self._row_buffer_raw: Any = None
         self._row_buffer_rows: List[Row] = []
         self._row_buffer_pos = 0
+        # Set once a fetchone() conversion failure has been logged at WARNING.
+        self._fetchone_conversion_warned = False
 
     def __iter__(self):
         while True:
@@ -219,7 +221,10 @@ def _raw_row_count(raw: Any) -> Optional[int]:
 
     @staticmethod
     def _convert_rows_individually(
-        raw: Any, convert: Callable[[Any], List[Row]], batch_error: Exception
+        raw: Any,
+        convert: Callable[[Any], List[Row]],
+        batch_error: Exception,
+        log_level: int = logging.WARNING,
     ) -> Optional[List[Any]]:
         """Convert ``raw`` row by row after a batch conversion failed.
 
@@ -228,12 +233,13 @@ def _convert_rows_individually(
         ``_FETCHONE_FALLBACK_PROBE_ROWS`` rows all fail with the same error as the
         batch (same type and message), the failure is treated as systematic and the
         remaining rows are parked with ``batch_error`` instead of being converted one
-        by one. Otherwise every row is tried individually. Returns None if ``raw`` has no rows to split, in
-        which case the caller re-raises the batch error.
+        by one. Otherwise every row is tried individually. Returns None if ``raw``
+        has no rows to split, in which case the caller re-raises the batch error.
         """
         n = ResultSet._raw_row_count(raw)
         if not n:
             return None
+
         def same_as_batch(r: Any) -> bool:
             return (
                 isinstance(r, _FailedRow)
@@ -247,7 +253,9 @@ def same_as_batch(r: Any) -> bool:
             # Only skip the remaining per-row retries when every probe row failed
             # with the same error as the batch; leading rows that fail for their
             # own reasons must not cause later good rows to be dropped.
-            if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(same_as_batch(r) for r in rows):
+            if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(
+                same_as_batch(r) for r in rows
+            ):
                 systematic = True
                 rows.extend(_FailedRow(batch_error) for _ in range(n - i))
                 break
@@ -259,14 +267,19 @@ def same_as_batch(r: Any) -> bool:
         failed = sum(isinstance(r, _FailedRow) for r in rows)
         # One summary line, no traceback: the error itself still reaches the caller
         # when the failing row is fetched. The full traceback is only logged at DEBUG.
-        logger.warning(
+        # ``log_level`` lets the caller log this at WARNING only once per result set.
+        logger.log(
+            log_level,
             "fetchone(): failed to convert %d of %d row(s) in a batch%s; each failing "
             "row raises when fetched (%s: %s)",
             failed,
             n,
-            " (first %d rows all failed, rest not retried)" % _FETCHONE_FALLBACK_PROBE_ROWS
-            if systematic
-            else "",
+            (
+                " (first %d rows all failed, rest not retried)"
+                % _FETCHONE_FALLBACK_PROBE_ROWS
+                if systematic
+                else ""
+            ),
             type(batch_error).__name__,
             batch_error,
         )
@@ -330,9 +343,19 @@ def _fetchone_buffered(
                 # as before batching, only the row(s) that fail to convert are
                 # lost: each failing row is parked as a _FailedRow and re-raises
                 # its own error when it is reached.
-                fallback = self._convert_rows_individually(raw, convert, batch_error)
+                # Warn once per result set; later failing batches (e.g. a column
+                # type that fails in every row) log at DEBUG to avoid a WARNING
+                # per refill for callers that catch and keep fetching.
+                warned = getattr(self, "_fetchone_conversion_warned", False)
+                fallback = self._convert_rows_individually(
+                    raw,
+                    convert,
+                    batch_error,
+                    log_level=logging.DEBUG if warned else logging.WARNING,
+                )
                 if fallback is None:
                     raise
+                self._fetchone_conversion_warned = True
                 rows = fallback
             # Check for drift before the empty-result early return, so a converter
             # that yields [] for a non-empty chunk raises instead of silently
@@ -555,7 +578,9 @@ def fetchmany_columnar(self, size: int):
             raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
         buffered = self._row_buffer_take_raw(size)
         remaining = size - (buffered.num_rows if buffered is not None else 0)
-        rest = self._fetchmany_columnar_from_queue(remaining) if remaining > 0 else None
+        rest = (
+            self._fetchmany_columnar_from_queue(remaining) if remaining > 0 else None
+        )
         if buffered is None and rest is None:
             rest = self._fetchmany_columnar_from_queue(0)
         return self._prepend_buffered(buffered, rest)
@@ -647,7 +672,9 @@ def _fetch_available_from_queue(self, size: int):
     def fetchall_arrow(self) -> "pyarrow.Table":
         """Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
         buffered = self._row_buffer_take_raw()
-        result_table = self._prepend_buffered(buffered, self._fetchall_arrow_from_queue())
+        result_table = self._prepend_buffered(
+            buffered, self._fetchall_arrow_from_queue()
+        )
         # If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table
         # Valid only for metadata commands result set
         if isinstance(result_table, ColumnTable) and pyarrow:
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 0d8b7e7bd..0236b17c7 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -601,6 +601,29 @@ def test_fetchone_convert_failure_on_every_row_reraises_per_row(self):
         self.assertEqual(rs._row_buffer_count(), 0)
         self.assertEqual(rs.fetchone()[0], 4)
 
+    def test_fetchone_convert_failure_warns_once_per_result_set(self):
+        import logging
+        from unittest.mock import patch
+
+        rs = self.make_columnar_result_set(12, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            side_effect=ValueError("bad value"),
+        ), self.assertLogs("databricks.sql.result_set", level="DEBUG") as logs:
+            # Three refill batches, every row failing the same way.
+            for _ in range(12):
+                with self.assertRaises(ValueError):
+                    rs.fetchone()
+        summaries = [
+            r for r in logs.records if "failed to convert" in r.getMessage()
+        ]
+        self.assertEqual(len(summaries), 3)
+        self.assertEqual(
+            [r.levelno for r in summaries],
+            [logging.WARNING, logging.DEBUG, logging.DEBUG],
+        )
+
     def test_fetchone_systematic_convert_failure_stops_per_row_retries(self):
         from unittest.mock import patch
         from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS

From 86b92d3dc6f6c1b476bea42db01efbbc6f4b40c9 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Thu, 29 Oct 2026 23:52:34 +0000
Subject: [PATCH 17/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4419073409 at src/databricks/sql/result_set.py:216
  - #4419073413 at tests/e2e/test_driver.py:997

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 32 +++++++++++++++-------
 tests/e2e/test_driver.py         | 14 +++++++---
 tests/e2e/test_kernel_backend.py | 12 ++++++---
 tests/unit/test_fetches.py       | 46 ++++++++++++++++++++++++++++++++
 4 files changed, 88 insertions(+), 16 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 60aab5eb3..4f07ada2b 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -196,20 +196,34 @@ def _row_buffer_advance(self, n: Optional[int]) -> Tuple[int, int]:
         return start, length
 
     def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
-        """Pop up to ``n`` (all if None) converted rows from the buffer."""
+        """Pop up to ``n`` (all if None) converted rows from the buffer.
+
+        If the requested range contains a ``_FailedRow``, only the rows up to and
+        including the first failed row are consumed and its error is raised. Rows
+        after it stay in the buffer for the next fetch. The good rows before it in
+        the same call are consumed with the failed fetch (a fetch either returns
+        all of its rows or raises), so ``rownumber`` stays aligned with what is
+        left in the buffer.
+        """
         if self._row_buffer_count() == 0:
             return []
         rows = self._row_buffer_rows
-        start, length = self._row_buffer_advance(n)
+        available = self._row_buffer_count()
+        length = available if n is None else max(0, min(n, available))
+        start = self._row_buffer_pos
+        failed: Optional[_FailedRow] = None
+        for i in range(start, start + length):
+            row = rows[i]
+            if isinstance(row, _FailedRow):
+                failed = row
+                length = i - start + 1
+                break
+        self._row_buffer_advance(length)
         if self._row_buffer_count() == 0:
             self._row_buffer_clear()
-        taken = rows[start : start + length]
-        for row in taken:
-            if isinstance(row, _FailedRow):
-                # The taken rows are consumed either way, matching how a failed
-                # conversion behaved before fetchone() buffering.
-                raise row.error.with_traceback(row.traceback)
-        return taken
+        if failed is not None:
+            raise failed.error.with_traceback(failed.traceback)
+        return rows[start : start + length]
 
     @staticmethod
     def _raw_row_count(raw: Any) -> Optional[int]:
diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py
index 47aa85979..5c645af6c 100644
--- a/tests/e2e/test_driver.py
+++ b/tests/e2e/test_driver.py
@@ -992,10 +992,16 @@ def _assert_rows_and_bound(self, rows, calls, arraysize):
         assert len(rows) == n
         assert [r[0] for r in rows] == list(range(n))
         assert all(r[1] == str(r[0]) for r in rows)
-        # Refills are capped per server batch / CloudFetch chunk, so an exact
-        # ceil(n/arraysize) bound is layout-dependent. Only distinguish "a few"
-        # conversions from "~n" (one per row, issue #981).
-        bound = n // 10
+        # Each refill converts up to min(arraysize, _FETCHONE_BUFFER_MAX_ROWS)
+        # rows, but refills are also capped at a server batch / CloudFetch chunk
+        # boundary, which can add one short refill per boundary. Allow 2x the
+        # ideal refill count plus a small constant for that slack — tight
+        # enough to catch tiny refills, not just one conversion per row (#981).
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        per_refill = max(min(arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
+        ideal = -(-n // per_refill)
+        bound = 2 * ideal + 10
         assert calls <= bound, (
             "row conversion called {} times for {} rows (arraysize={}); "
             "expected at most {} (issue #981)".format(calls, n, arraysize, bound)
diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py
index e1259af0e..e5132644d 100644
--- a/tests/e2e/test_kernel_backend.py
+++ b/tests/e2e/test_kernel_backend.py
@@ -206,9 +206,15 @@ def _assert_rows_and_bound(rows, calls):
     assert len(rows) == _ROW_BUFFER_N
     assert [r[0] for r in rows] == list(range(_ROW_BUFFER_N))
     assert all(r[1] == str(r[0]) for r in rows)
-    # Server batch boundaries add short refills, so an exact ceil(N/arraysize)
-    # bound is layout-dependent. Only distinguish "a few" from "~N" (issue #981).
-    bound = _ROW_BUFFER_N // 10
+    # Each refill converts up to min(arraysize, _FETCHONE_BUFFER_MAX_ROWS) rows;
+    # server batch boundaries can add one short refill each. Allow 2x the ideal
+    # refill count plus a small constant — tight enough to catch tiny refills,
+    # not just one conversion per row (issue #981).
+    from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+    per_refill = max(min(_ROW_BUFFER_ARRAYSIZE, _FETCHONE_BUFFER_MAX_ROWS), 1)
+    ideal = -(-_ROW_BUFFER_N // per_refill)
+    bound = 2 * ideal + 10
     assert calls <= bound, (
         f"_convert_arrow_table called {calls} times for {_ROW_BUFFER_N} rows; "
         f"expected at most {bound} (issue #981)"
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 0236b17c7..2bda58a99 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -585,6 +585,52 @@ def convert_failing_on_2(self_, table):
         self.assertTrue(any("1 of 4 row(s)" in m for m in logs.output))
         self.assertEqual(rs.rownumber, 10)
 
+    def _patch_convert_failing_on_2(self):
+        from unittest.mock import patch
+
+        real_convert = ThriftResultSet._convert_columnar_table
+
+        def convert_failing_on_2(self_, table):
+            rows = real_convert(self_, table)
+            if any(r[0] == 2 for r in rows):
+                raise ValueError("bad value")
+            return rows
+
+        return patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert_failing_on_2,
+        )
+
+    def test_fetchmany_over_failed_buffered_row_keeps_rows_after_it(self):
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with self._patch_convert_failing_on_2(), self.assertLogs(
+            "databricks.sql.result_set", level="WARNING"
+        ):
+            # Buffer now holds rows 1, 2 (failed), 3.
+            self.assertEqual(rs.fetchone()[0], 0)
+            with self.assertRaises(ValueError):
+                rs.fetchmany(3)
+            # Consumed up to and including the failed row; row 3 stays buffered.
+            self.assertEqual(rs.rownumber, 3)
+            self.assertEqual(rs._row_buffer_count(), 1)
+            self.assertEqual([r[0] for r in rs.fetchmany(3)], [3, 4, 5])
+            self.assertEqual([r[0] for r in rs.fetchall()], [6, 7, 8, 9])
+        self.assertEqual(rs.rownumber, 10)
+
+    def test_fetchall_over_failed_buffered_row_keeps_rows_after_it(self):
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with self._patch_convert_failing_on_2(), self.assertLogs(
+            "databricks.sql.result_set", level="WARNING"
+        ):
+            self.assertEqual(rs.fetchone()[0], 0)
+            with self.assertRaises(ValueError):
+                rs.fetchall()
+            self.assertEqual(rs.rownumber, 3)
+            self.assertEqual([r[0] for r in rs.fetchall()], list(range(3, 10)))
+        self.assertEqual(rs.rownumber, 10)
+
     def test_fetchone_convert_failure_on_every_row_reraises_per_row(self):
         from unittest.mock import patch
 

From c848f81d4649bb48db46e09254287143e04e4e82 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:07:24 +0000
Subject: [PATCH 18/29] ai: apply changes for #1007 (3 review threads)

Addresses:
  - #4419116506 at src/databricks/sql/result_set.py:35
  - #4419116511 at src/databricks/sql/result_set.py:448
  - #4419116514 at src/databricks/sql/result_set.py:358

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py  | 41 +++++++++++++++---
 tests/unit/test_fetches.py        | 72 ++++++++++++++++++++++++++++++-
 tests/unit/test_sea_result_set.py | 27 ++++++++++++
 3 files changed, 132 insertions(+), 8 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 4f07ada2b..238d2cec9 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -29,11 +29,18 @@
 logger = logging.getLogger(__name__)
 
 # Upper bound on how many rows a fetchone() refill pulls and converts at once. The
-# refill uses ``arraysize`` but is capped here so the first fetchone() on a cursor
-# with a large arraysize (default 100000) doesn't materialise a full page just to
-# return one row.
+# refill uses ``arraysize`` but is capped here so a cursor with a large arraysize
+# (default 100000) doesn't materialise a full page just to return one row.
 _FETCHONE_BUFFER_MAX_ROWS = 1000
 
+# fetchone() refills grow geometrically from this size (1 -> 10 -> 100 -> ...) up to
+# min(arraysize, _FETCHONE_BUFFER_MAX_ROWS). The first refill converts a single row,
+# so a lone fetchone() (e.g. SQLAlchemy ``.first()`` / ``.scalar()`` or an EXISTS
+# probe) costs the same as before batching, while iteration loops quickly reach the
+# full batch size.
+_FETCHONE_BUFFER_INITIAL_ROWS = 1
+_FETCHONE_BUFFER_GROWTH_FACTOR = 10
+
 # After a fetchone() batch conversion fails, the batch is retried one row at a time.
 # If this many leading rows all fail too, with the same error type and message as
 # the batch, the failure is treated as systematic: the rest of the batch is not
@@ -124,6 +131,9 @@ def __init__(
         self._row_buffer_pos = 0
         # Set once a fetchone() conversion failure has been logged at WARNING.
         self._fetchone_conversion_warned = False
+        # Size of the next fetchone() refill; grows on each refill (see
+        # ``_FETCHONE_BUFFER_GROWTH_FACTOR``).
+        self._fetchone_refill_rows = _FETCHONE_BUFFER_INITIAL_ROWS
 
     def __iter__(self):
         while True:
@@ -336,8 +346,11 @@ def _fetchone_buffered(
     ) -> Optional[Row]:
         """Serve fetchone() from the row buffer, refilling in small batches.
 
-        Each refill asks ``fetch_raw`` for up to ``arraysize`` rows, capped at
-        ``_FETCHONE_BUFFER_MAX_ROWS``. ``fetch_raw`` may return fewer (Thrift and the
+        Refill sizes start at ``_FETCHONE_BUFFER_INITIAL_ROWS`` and grow by
+        ``_FETCHONE_BUFFER_GROWTH_FACTOR`` on each refill, capped at ``arraysize``
+        and ``_FETCHONE_BUFFER_MAX_ROWS``, so a single fetchone() converts only one
+        row while iteration settles into full batches. ``fetch_raw`` may return
+        fewer (Thrift and the
         kernel backend only serve what is already buffered, refilling at most until
         one row is available, so latency to the first row is unchanged).
 
@@ -346,9 +359,22 @@ def _fetchone_buffered(
         queue is only read again once the buffer is empty, so at that point the
         index equals the queue position again (Thrift relies on that for
         ``expected_row_start_offset``).
+
+        If conversion yields a different number of rows than the raw chunk held,
+        ``InternalError`` is raised and that whole chunk is dropped (the queue and
+        ``rownumber`` have already moved past it). The result set should be
+        considered unusable after that error.
         """
         if self._row_buffer_count() == 0:
-            raw = fetch_raw(max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1))
+            cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
+            target = getattr(
+                self, "_fetchone_refill_rows", _FETCHONE_BUFFER_INITIAL_ROWS
+            )
+            size = max(min(target, cap), 1)
+            self._fetchone_refill_rows = min(
+                size * _FETCHONE_BUFFER_GROWTH_FACTOR, _FETCHONE_BUFFER_MAX_ROWS
+            )
+            raw = fetch_raw(size)
             try:
                 rows = convert(raw)
             except Exception as batch_error:
@@ -379,6 +405,9 @@ def _fetchone_buffered(
                 # fetch_raw advanced _next_row_index (and the Thrift fetch offset)
                 # by raw_count; the buffer must line up 1:1 with the raw chunk or
                 # rownumber and fetchmany_arrow slicing would silently drift.
+                # The raw chunk is not buffered, so its raw_count rows are lost
+                # (yet counted in rownumber): treat the result set as unusable
+                # after this error; later fetches resume from the next chunk.
                 raise InternalError(
                     "fetchone(): converted {} row(s) from a chunk of {} raw row(s)".format(
                         len(rows), raw_count
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 2bda58a99..9c6c7cef3 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -370,6 +370,25 @@ def test_mixed_fetches_across_pages_keep_order_and_refill_offsets(self):
         self.assertEqual(rs.rownumber, 10)
         self.assertEqual(offsets, [0, 3, 6, 9])
 
+    def test_close_drops_fetchone_row_buffer(self):
+        # Rows parked in the fetchone() buffer must not survive close(). The warm
+        # refill pulls every row into the buffer, so the queue itself is empty and
+        # any row returned after close() could only come from a stale buffer.
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        for fetch_after_close in (
+            lambda rs: rs.fetchall(),
+            lambda rs: rs.fetchmany_arrow(10).to_pylist(),
+        ):
+            rs = self.make_dummy_result_set_from_initial_results([[1], [2], [3], [4]])
+            rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
+            self.assertSequenceEqual(rs.fetchone(), [1])
+            self.assertEqual(rs._row_buffer_count(), 3)
+            rs.close()
+            self.assertEqual(rs._row_buffer_count(), 0)
+            self.assertEqual(fetch_after_close(rs), [])
+            self.assertIsNone(rs.fetchone())
+
     def test_fetchone_refill_does_not_fetch_extra_pages(self):
         # A refill must not make extra round-trips to fill ``arraysize`` when
         # server pages are smaller (e.g. byte-limited): it serves what the current
@@ -408,6 +427,10 @@ def test_cloudfetch_fetchone_refill_capped_to_current_chunk(self):
             description=[("col0", "integer", None, None, None, None, None)],
         )
         rs.arraysize = 100
+        # Skip the adaptive refill ramp so the chunk cap, not the ramp, bounds it.
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
 
         self.assertSequenceEqual(rs.fetchone(), [8])
         queue.next_n_rows.assert_called_once_with(2)
@@ -482,13 +505,21 @@ class ColumnarFetchoneBufferTests(unittest.TestCase):
     rows in batches (not one row at a time) while keeping order and rownumber."""
 
     @staticmethod
-    def make_columnar_result_set(n_rows, arraysize):
+    def make_columnar_result_set(n_rows, arraysize, warm=True):
+        """Build a result set over ``n_rows`` columnar rows.
+
+        With ``warm`` (the default) the adaptive fetchone() refill ramp is skipped,
+        so every refill is ``min(arraysize, _FETCHONE_BUFFER_MAX_ROWS)`` rows and the
+        batch-boundary assertions below stay simple. Pass ``warm=False`` to test
+        the ramp itself.
+        """
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
         from databricks.sql.utils import ColumnQueue, ColumnTable
 
         table = ColumnTable([list(range(n_rows))], ["col0"])
         mock_thrift_backend = Mock(spec=ThriftDatabricksClient)
         mock_thrift_backend.fetch_results.return_value = (ColumnQueue(table), False, 0)
-        return ThriftResultSet(
+        rs = ThriftResultSet(
             connection=Mock(),
             execute_response=ExecuteResponse(
                 command_id=None,
@@ -501,6 +532,9 @@ def make_columnar_result_set(n_rows, arraysize):
             thrift_client=mock_thrift_backend,
             arraysize=arraysize,
         )
+        if warm:
+            rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
+        return rs
 
     @staticmethod
     def _count_conversions():
@@ -543,6 +577,40 @@ def test_fetchone_refill_is_capped_for_large_arraysize(self):
         self.assertEqual([r[0] for r in rs.fetchall()], list(range(1, n_rows)))
         self.assertEqual(rs.rownumber, n_rows)
 
+    def test_first_fetchone_converts_a_single_row(self):
+        # A lone fetchone() (SQLAlchemy .first()/.scalar(), EXISTS probes) must not
+        # pull and convert a whole batch: the first refill is a single row.
+        rs = self.make_columnar_result_set(5000, arraysize=100000, warm=False)
+        with self._count_conversions() as m:
+            self.assertEqual(rs.fetchone()[0], 0)
+        self.assertEqual(m.call_count, 1)
+        self.assertEqual(m.call_args[0][1].num_rows, 1)
+        self.assertEqual(rs._row_buffer_count(), 0)
+        self.assertEqual(rs.rownumber, 1)
+
+    def test_fetchone_refill_size_ramps_up_to_cap(self):
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        n_rows = 1 + 10 + 100 + _FETCHONE_BUFFER_MAX_ROWS * 2
+        rs = self.make_columnar_result_set(n_rows, arraysize=100000, warm=False)
+        with self._count_conversions() as m:
+            rows = [r[0] for r in rs]
+        self.assertEqual(rows, list(range(n_rows)))
+        sizes = [c[0][1].num_rows for c in m.call_args_list]
+        self.assertEqual(
+            sizes[:5],
+            [1, 10, 100, _FETCHONE_BUFFER_MAX_ROWS, _FETCHONE_BUFFER_MAX_ROWS],
+        )
+        self.assertEqual(rs.rownumber, n_rows)
+
+    def test_fetchone_refill_ramp_is_capped_by_arraysize(self):
+        rs = self.make_columnar_result_set(10, arraysize=4, warm=False)
+        with self._count_conversions() as m:
+            rows = [r[0] for r in rs]
+        self.assertEqual(rows, list(range(10)))
+        sizes = [c[0][1].num_rows for c in m.call_args_list]
+        self.assertEqual(sizes[:4], [1, 4, 4, 1])
+
     def test_mixed_fetches_keep_order_and_rownumber(self):
         rs = self.make_columnar_result_set(10, arraysize=4)
         self.assertEqual(rs.fetchone()[0], 0)
diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py
index f78ee2ba9..db65c4a77 100644
--- a/tests/unit/test_sea_result_set.py
+++ b/tests/unit/test_sea_result_set.py
@@ -308,6 +308,29 @@ def test_close_when_connection_closed(
         assert result_set.has_been_closed_server_side is True
         assert result_set.status == CommandState.CLOSED
 
+    @pytest.mark.parametrize("fetch_after_close", ["fetchall", "fetchmany_arrow"])
+    def test_close_drops_fetchone_row_buffer(
+        self, result_set_with_data, fetch_after_close
+    ):
+        """Rows parked in the fetchone() buffer must not survive close()."""
+        if fetch_after_close == "fetchmany_arrow" and pyarrow is None:
+            pytest.skip("PyArrow is not installed")
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        rs = result_set_with_data
+        # Warm refill pulls all 5 rows into the buffer, emptying the queue, so any
+        # row returned after close() could only come from a stale buffer.
+        rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
+        assert rs.fetchone()[0] == "value1"
+        assert rs._row_buffer_count() == 4
+        rs.close()
+        assert rs._row_buffer_count() == 0
+        if fetch_after_close == "fetchall":
+            assert rs.fetchall() == []
+        else:
+            assert rs.fetchmany_arrow(10).num_rows == 0
+        assert rs.fetchone() is None
+
     def test_convert_json_types(self, result_set_with_data, sample_data):
         """Test the _convert_json_types method."""
         # Call _convert_json_types
@@ -690,6 +713,10 @@ def test_cloudfetch_fetchone_refill_capped_to_current_chunk(
                 buffer_size_bytes=1000,
                 arraysize=100,
             )
+        # Skip the adaptive refill ramp so the chunk cap, not the ramp, bounds it.
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
         assert rs.fetchone().col2 == 8
         queue.next_n_rows.assert_called_once_with(2)
         assert rs.fetchone().col2 == 9

From 77e1090555121f353cbaf33f885faeacfebe0569 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:17:31 +0000
Subject: [PATCH 19/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4419207269 at src/databricks/sql/result_set.py:280

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 83 ++++++++++++--------------------
 tests/unit/test_fetches.py       | 63 +++++++++++++++++-------
 2 files changed, 77 insertions(+), 69 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 238d2cec9..19fda80e9 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -41,15 +41,6 @@
 _FETCHONE_BUFFER_INITIAL_ROWS = 1
 _FETCHONE_BUFFER_GROWTH_FACTOR = 10
 
-# After a fetchone() batch conversion fails, the batch is retried one row at a time.
-# If this many leading rows all fail too, with the same error type and message as
-# the batch, the failure is treated as systematic: the rest of the batch is not
-# retried row by row (each would cost a separate conversion) and those rows re-raise
-# the original batch error instead. Leading rows failing for other reasons do not
-# trigger the shortcut, so later good rows are still converted.
-_FETCHONE_FALLBACK_PROBE_ROWS = 5
-
-
 class _FailedRow:
     """Placeholder in the fetchone() row buffer for a row that failed to convert."""
 
@@ -57,8 +48,8 @@ class _FailedRow:
 
     def __init__(self, error: Exception):
         self.error = error
-        # Snapshot the conversion traceback. Systematic failures park the same
-        # exception instance in many rows; restoring this before each raise keeps
+        # Snapshot the conversion traceback. A converter may raise the same
+        # exception instance for many rows; restoring this before each raise keeps
         # its __traceback__ from growing (and pinning frames) on every fetchone().
         self.traceback = error.__traceback__
 
@@ -250,60 +241,50 @@ def _convert_rows_individually(
         batch_error: Exception,
         log_level: int = logging.WARNING,
     ) -> Optional[List[Any]]:
-        """Convert ``raw`` row by row after a batch conversion failed.
-
-        Rows that fail are returned as ``_FailedRow`` placeholders so the error is
-        raised only when that row is handed out. If the first
-        ``_FETCHONE_FALLBACK_PROBE_ROWS`` rows all fail with the same error as the
-        batch (same type and message), the failure is treated as systematic and the
-        remaining rows are parked with ``batch_error`` instead of being converted one
-        by one. Otherwise every row is tried individually. Returns None if ``raw``
-        has no rows to split, in which case the caller re-raises the batch error.
+        """Convert ``raw`` in smaller pieces after a batch conversion failed.
+
+        The failed batch is bisected: each failing range is split in half and the
+        halves are converted separately, down to single rows. Only rows that fail
+        on their own are returned as ``_FailedRow`` placeholders (raised when that
+        row is handed out); every row that converts is kept, so a failure is never
+        attributed to a row that was not itself tried. Cost is about
+        ``2 * k * log2(n)`` conversions for ``k`` bad rows, and at most ``2n``
+        when every row fails. Returns None if ``raw`` has no rows to split, in
+        which case the caller re-raises the batch error.
         """
         n = ResultSet._raw_row_count(raw)
         if not n:
             return None
 
-        def same_as_batch(r: Any) -> bool:
-            return (
-                isinstance(r, _FailedRow)
-                and type(r.error) is type(batch_error)
-                and str(r.error) == str(batch_error)
-            )
-
         rows: List[Any] = []
-        systematic = False
-        for i in range(n):
-            # Only skip the remaining per-row retries when every probe row failed
-            # with the same error as the batch; leading rows that fail for their
-            # own reasons must not cause later good rows to be dropped.
-            if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(
-                same_as_batch(r) for r in rows
-            ):
-                systematic = True
-                rows.extend(_FailedRow(batch_error) for _ in range(n - i))
-                break
-            one = raw[i : i + 1] if isinstance(raw, list) else raw.slice(i, 1)
-            try:
-                rows.extend(convert(one))
-            except Exception as e:
-                rows.append(_FailedRow(e))
+
+        def convert_range(lo: int, hi: int, error: Optional[Exception]) -> None:
+            # ``error`` is the failure already seen for [lo, hi), if it was tried.
+            if error is None:
+                piece = raw[lo:hi] if isinstance(raw, list) else raw.slice(lo, hi - lo)
+                try:
+                    rows.extend(convert(piece))
+                    return
+                except Exception as e:
+                    error = e
+            if hi - lo == 1:
+                rows.append(_FailedRow(error))
+                return
+            mid = (lo + hi) // 2
+            convert_range(lo, mid, None)
+            convert_range(mid, hi, None)
+
+        convert_range(0, n, batch_error)
         failed = sum(isinstance(r, _FailedRow) for r in rows)
         # One summary line, no traceback: the error itself still reaches the caller
         # when the failing row is fetched. The full traceback is only logged at DEBUG.
         # ``log_level`` lets the caller log this at WARNING only once per result set.
         logger.log(
             log_level,
-            "fetchone(): failed to convert %d of %d row(s) in a batch%s; each failing "
+            "fetchone(): failed to convert %d of %d row(s) in a batch; each failing "
             "row raises when fetched (%s: %s)",
             failed,
             n,
-            (
-                " (first %d rows all failed, rest not retried)"
-                % _FETCHONE_FALLBACK_PROBE_ROWS
-                if systematic
-                else ""
-            ),
             type(batch_error).__name__,
             batch_error,
         )
@@ -379,7 +360,7 @@ def _fetchone_buffered(
                 rows = convert(raw)
             except Exception as batch_error:
                 # The queue has already moved past this batch, so it can't be
-                # re-read. Fall back to converting it one row at a time so that,
+                # re-read. Fall back to bisecting it down to single rows so that,
                 # as before batching, only the row(s) that fail to convert are
                 # lost: each failing row is parked as a _FailedRow and re-raises
                 # its own error when it is reached.
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 9c6c7cef3..d25fb634a 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -738,19 +738,16 @@ def test_fetchone_convert_failure_warns_once_per_result_set(self):
             [logging.WARNING, logging.DEBUG, logging.DEBUG],
         )
 
-    def test_fetchone_systematic_convert_failure_stops_per_row_retries(self):
+    def test_fetchone_convert_failure_on_every_row_bisects_to_single_rows(self):
         from unittest.mock import patch
-        from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS
 
-        batch_error = ValueError("batch failure")
+        shared_error = ValueError("batch failure")
         calls = []
 
         def convert(self_, table):
             calls.append(table.num_rows)
-            if len(calls) == 1:
-                raise batch_error
-            # Same type and message as the batch error: a systematic failure.
-            raise ValueError("batch failure")
+            # The converter re-raises one shared exception instance for every row.
+            raise shared_error
 
         rs = self.make_columnar_result_set(50, arraysize=20)
         with patch.object(
@@ -763,27 +760,57 @@ def convert(self_, table):
             for i in range(20):
                 with self.assertRaises(ValueError) as ctx:
                     rs.fetchone()
-                if i >= _FETCHONE_FALLBACK_PROBE_ROWS:
-                    self.assertIs(ctx.exception, batch_error)
-                    tb, depth = ctx.exception.__traceback__, 0
-                    while tb is not None:
-                        tb, depth = tb.tb_next, depth + 1
-                    tb_depths.append(depth)
+                self.assertIs(ctx.exception, shared_error)
+                tb, depth = ctx.exception.__traceback__, 0
+                while tb is not None:
+                    tb, depth = tb.tb_next, depth + 1
+                tb_depths.append(depth)
                 self.assertEqual(rs.rownumber, i + 1)
-            # Re-raising the shared batch error must not grow its traceback per row.
+            # Re-raising a shared error must not grow its traceback per row.
             self.assertEqual(len(set(tb_depths)), 1)
-        # One batch attempt plus only the probe rows were converted individually.
-        self.assertEqual(len(calls), 1 + _FETCHONE_FALLBACK_PROBE_ROWS)
+        # Bisection of an all-failing batch of 20 tries every node of the tree once:
+        # 2 * 20 - 1 conversions (including the initial batch attempt).
+        self.assertEqual(calls[0], 20)
+        self.assertEqual(len(calls), 2 * 20 - 1)
+        self.assertEqual(sum(1 for c in calls if c == 1), 20)
         self.assertTrue(any("20 of 20 row(s)" in m for m in logs.output))
         self.assertTrue(any("batch failure" in m for m in logs.output))
         self.assertEqual(rs.fetchone()[0], 20)
 
+    def test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows(self):
+        from unittest.mock import patch
+
+        real_convert = ThriftResultSet._convert_columnar_table
+        # Enough leading bad rows, all failing with the same message as the batch
+        # (e.g. "year 0 is out of range"), to look like a systematic failure.
+        n_bad = 8
+
+        def convert(self_, table):
+            rows = real_convert(self_, table)
+            if any(r[0] < n_bad for r in rows):
+                raise ValueError("year 0 is out of range")
+            return rows
+
+        rs = self.make_columnar_result_set(20, arraysize=20)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert,
+        ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
+            for _ in range(n_bad):
+                with self.assertRaises(ValueError):
+                    rs.fetchone()
+            # Every row after the bad run still converts and is served.
+            self.assertEqual([r[0] for r in rs.fetchall()], list(range(n_bad, 20)))
+        self.assertTrue(any("{} of 20 row(s)".format(n_bad) in m for m in logs.output))
+        self.assertEqual(rs.rownumber, 20)
+
     def test_fetchone_leading_bad_rows_do_not_drop_later_good_rows(self):
         from unittest.mock import patch
-        from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS
 
         real_convert = ThriftResultSet._convert_columnar_table
-        n_bad = _FETCHONE_FALLBACK_PROBE_ROWS + 1
+        n_bad = 6
 
         def convert(self_, table):
             rows = real_convert(self_, table)

From 46012403cd74ea7514c562b058eb2fd450b64d03 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:27:47 +0000
Subject: [PATCH 20/29] ai: apply changes for #1007 (3 review threads)

Addresses:
  - #4419261192 at src/databricks/sql/result_set.py:44
  - #4419261198 at src/databricks/sql/result_set.py:265
  - #4419261206 at src/databricks/sql/backend/sea/result_set.py:280

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 .../sql/backend/kernel/result_set.py          | 16 +++----
 src/databricks/sql/backend/sea/result_set.py  | 25 +++++-----
 src/databricks/sql/result_set.py              | 46 +++++++++++++------
 tests/unit/test_fetches.py                    | 37 +++++++++++++++
 4 files changed, 89 insertions(+), 35 deletions(-)

diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index d269d6454..04754041d 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -245,17 +245,15 @@ def fetchone(self) -> Optional[Row]:
     def fetchmany(self, size: int) -> List[Row]:
         if size < 0:
             raise ValueError(f"fetchmany size must be >= 0, got {size}")
-        buffered = self._row_buffer_take_rows(size)
-        remaining = size - len(buffered)
-        if remaining == 0:
-            return buffered
-        rest = self._convert_arrow_table(self._fetchmany_arrow_unbuffered(remaining))
-        return buffered + rest if buffered else rest
+        return self._fetch_rows_with_buffer(
+            size,
+            lambda n: self._convert_arrow_table(self._fetchmany_arrow_unbuffered(n)),
+        )
 
     def fetchall(self) -> List[Row]:
-        buffered = self._row_buffer_take_rows()
-        rest = self._convert_arrow_table(self._drain())
-        return buffered + rest if buffered else rest
+        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.
diff --git a/src/databricks/sql/backend/sea/result_set.py b/src/databricks/sql/backend/sea/result_set.py
index 09b53e842..38e4727c2 100644
--- a/src/databricks/sql/backend/sea/result_set.py
+++ b/src/databricks/sql/backend/sea/result_set.py
@@ -279,15 +279,13 @@ def fetchmany(self, size: int) -> List[Row]:
 
         if size < 0:
             raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
-        buffered = self._row_buffer_take_rows(size)
-        remaining = size - len(buffered)
-        if buffered and remaining == 0:
-            return buffered
         if isinstance(self.results, JsonQueue):
-            rest = self._create_json_table(self.fetchmany_json(remaining))
-        else:
-            rest = self._convert_arrow_table(self.fetchmany_arrow(remaining))
-        return buffered + rest if buffered else rest
+            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]:
         """
@@ -297,9 +295,10 @@ def fetchall(self) -> List[Row]:
             List of Row objects containing all remaining rows
         """
 
-        buffered = self._row_buffer_take_rows()
         if isinstance(self.results, JsonQueue):
-            rest = self._create_json_table(self.fetchall_json())
-        else:
-            rest = self._convert_arrow_table(self.fetchall_arrow())
-        return buffered + rest if buffered else rest
+            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())
+        )
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 19fda80e9..23ea4290f 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -41,6 +41,7 @@
 _FETCHONE_BUFFER_INITIAL_ROWS = 1
 _FETCHONE_BUFFER_GROWTH_FACTOR = 10
 
+
 class _FailedRow:
     """Placeholder in the fetchone() row buffer for a row that failed to convert."""
 
@@ -226,6 +227,26 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
             raise failed.error.with_traceback(failed.traceback)
         return rows[start : start + length]
 
+    def _fetch_rows_with_buffer(
+        self,
+        size: Optional[int],
+        fetch_rest: Callable[[Optional[int]], List[Row]],
+    ) -> List[Row]:
+        """Return ``size`` rows (all if None), draining the fetchone() buffer first.
+
+        Every row-level fetchmany()/fetchall() goes through here, so rows the
+        fetchone() buffer already holds are always returned before newly fetched
+        ones. ``fetch_rest(n)`` must fetch up to ``n`` (all if None) rows
+        while bypassing the buffer. It is not called when the buffer already
+        covers ``size``.
+        """
+        buffered = self._row_buffer_take_rows(size)
+        remaining = None if size is None else size - len(buffered)
+        if buffered and remaining == 0:
+            return buffered
+        rest = fetch_rest(remaining)
+        return buffered + rest if buffered else rest
+
     @staticmethod
     def _raw_row_count(raw: Any) -> Optional[int]:
         """Row count of a raw chunk (Arrow table, ColumnTable or list); None if unknown."""
@@ -774,12 +795,13 @@ def fetchall(self) -> List[Row]:
         """
         Fetch all (remaining) rows of a query result, returning them as a list of rows.
         """
-        buffered = self._row_buffer_take_rows()
         if isinstance(self.results, ColumnQueue):
-            rest = self._convert_columnar_table(self.fetchall_columnar())
-        else:
-            rest = self._convert_arrow_table(self.fetchall_arrow())
-        return buffered + rest if buffered else rest
+            return self._fetch_rows_with_buffer(
+                None, lambda _: self._convert_columnar_table(self.fetchall_columnar())
+            )
+        return self._fetch_rows_with_buffer(
+            None, lambda _: self._convert_arrow_table(self.fetchall_arrow())
+        )
 
     def fetchmany(self, size: int) -> List[Row]:
         """
@@ -789,15 +811,13 @@ def fetchmany(self, size: int) -> List[Row]:
         """
         if size < 0:
             raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
-        buffered = self._row_buffer_take_rows(size)
-        remaining = size - len(buffered)
-        if buffered and remaining == 0:
-            return buffered
         if isinstance(self.results, ColumnQueue):
-            rest = self._convert_columnar_table(self.fetchmany_columnar(remaining))
-        else:
-            rest = self._convert_arrow_table(self.fetchmany_arrow(remaining))
-        return buffered + rest if buffered else rest
+            return self._fetch_rows_with_buffer(
+                size, lambda n: self._convert_columnar_table(self.fetchmany_columnar(n))
+            )
+        return self._fetch_rows_with_buffer(
+            size, lambda n: self._convert_arrow_table(self.fetchmany_arrow(n))
+        )
 
     @staticmethod
     def _get_schema_description(table_schema_message):
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index d25fb634a..44653279f 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -777,6 +777,43 @@ def convert(self_, table):
         self.assertTrue(any("batch failure" in m for m in logs.output))
         self.assertEqual(rs.fetchone()[0], 20)
 
+    def test_fetchone_all_failing_max_batch_conversion_count_is_bounded(self):
+        """Document the worst-case cost of bisecting a full all-failing refill.
+
+        When every row fails (a systemic converter / type problem), bisection tries
+        every node of the tree: exactly ``2n - 1`` conversions for an ``n``-row
+        refill, i.e. about 2x the one-conversion-per-row cost of fetchone() before
+        batching. Short-circuiting on "same exception type in both halves" is
+        deliberately not done, since it would drop good rows that share the error
+        type (see test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows).
+        """
+        from unittest.mock import patch
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        n = _FETCHONE_BUFFER_MAX_ROWS
+        calls = []
+
+        def convert(self_, table):
+            calls.append(table.num_rows)
+            raise ValueError("systemic failure")
+
+        rs = self.make_columnar_result_set(n, arraysize=n)
+        # Skip the adaptive refill ramp so the first refill is a full batch.
+        rs._fetchone_refill_rows = n
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert,
+        ), self.assertLogs("databricks.sql.result_set", level="WARNING"):
+            with self.assertRaises(ValueError):
+                rs.fetchone()
+        self.assertEqual(calls[0], n)
+        self.assertEqual(len(calls), 2 * n - 1)
+        # Total rows converted across all attempts is O(n log n), not O(n^2).
+        self.assertLessEqual(sum(calls), n * (n - 1).bit_length() + n)
+        self.assertEqual(rs._row_buffer_count(), n - 1)
+
     def test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows(self):
         from unittest.mock import patch
 

From cfbc6f6fe099d202aa5b0770861ba05e6f423fac Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:37:07 +0000
Subject: [PATCH 21/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4419314538 at src/databricks/sql/result_set.py:395

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 14 +++++++++
 tests/unit/test_fetches.py       | 53 ++++++++++++++++++++++++++++++--
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 23ea4290f..bdba7e164 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -41,6 +41,10 @@
 _FETCHONE_BUFFER_INITIAL_ROWS = 1
 _FETCHONE_BUFFER_GROWTH_FACTOR = 10
 
+# Conversion errors that are not caused by row data. fetchone() re-raises these
+# straight away instead of bisecting the batch to find the failing row(s).
+_NON_DATA_CONVERSION_ERRORS = (MemoryError, RecursionError)
+
 
 class _FailedRow:
     """Placeholder in the fetchone() row buffer for a row that failed to convert."""
@@ -379,6 +383,10 @@ def _fetchone_buffered(
             raw = fetch_raw(size)
             try:
                 rows = convert(raw)
+            except _NON_DATA_CONVERSION_ERRORS:
+                # Resource / interpreter failures say nothing about individual
+                # rows; retrying smaller slices would only repeat (or worsen) them.
+                raise
             except Exception as batch_error:
                 # The queue has already moved past this batch, so it can't be
                 # re-read. Fall back to bisecting it down to single rows so that,
@@ -399,6 +407,12 @@ def _fetchone_buffered(
                     raise
                 self._fetchone_conversion_warned = True
                 rows = fallback
+                if all(isinstance(r, _FailedRow) for r in rows):
+                    # Every row failed: likely systemic (a column type or library
+                    # problem that hits each row). Restart the refill ramp so later
+                    # refills convert one row at a time, as before batching, rather
+                    # than paying a full ~2n-conversion bisection on every refill.
+                    self._fetchone_refill_rows = _FETCHONE_BUFFER_INITIAL_ROWS
             # Check for drift before the empty-result early return, so a converter
             # that yields [] for a non-empty chunk raises instead of silently
             # ending iteration and losing those rows.
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 44653279f..1d977f9e6 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -725,19 +725,66 @@ def test_fetchone_convert_failure_warns_once_per_result_set(self):
             "_convert_columnar_table",
             side_effect=ValueError("bad value"),
         ), self.assertLogs("databricks.sql.result_set", level="DEBUG") as logs:
-            # Three refill batches, every row failing the same way.
+            # One 4-row refill, then (ramp reset after an all-failing batch) eight
+            # 1-row refills, every row failing the same way.
             for _ in range(12):
                 with self.assertRaises(ValueError):
                     rs.fetchone()
         summaries = [
             r for r in logs.records if "failed to convert" in r.getMessage()
         ]
-        self.assertEqual(len(summaries), 3)
+        self.assertEqual(len(summaries), 9)
         self.assertEqual(
             [r.levelno for r in summaries],
-            [logging.WARNING, logging.DEBUG, logging.DEBUG],
+            [logging.WARNING] + [logging.DEBUG] * 8,
         )
 
+    def test_fetchone_all_failing_batch_resets_refill_ramp(self):
+        from unittest.mock import patch
+
+        calls = []
+
+        def convert(self_, table):
+            calls.append(table.num_rows)
+            raise ValueError("systemic failure")
+
+        rs = self.make_columnar_result_set(30, arraysize=10)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert,
+        ), self.assertLogs("databricks.sql.result_set", level="DEBUG"):
+            for _ in range(15):
+                with self.assertRaises(ValueError):
+                    rs.fetchone()
+        # First full refill is bisected (2 * 10 - 1 conversions); after that every
+        # refill is a single row costing exactly one conversion.
+        self.assertEqual(len(calls), 2 * 10 - 1 + 5)
+        self.assertEqual(calls[-5:], [1] * 5)
+        self.assertEqual(rs.rownumber, 15)
+
+    def test_fetchone_memory_error_is_not_bisected(self):
+        from unittest.mock import patch
+
+        calls = []
+
+        def convert(self_, table):
+            calls.append(table.num_rows)
+            raise MemoryError()
+
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert,
+        ):
+            with self.assertRaises(MemoryError):
+                rs.fetchone()
+        self.assertEqual(calls, [4])
+        self.assertEqual(rs._row_buffer_count(), 0)
+
     def test_fetchone_convert_failure_on_every_row_bisects_to_single_rows(self):
         from unittest.mock import patch
 

From 238dd78a3132294dfae458dfe2df64554ad17094 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:44:20 +0000
Subject: [PATCH 22/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4419356673 at src/databricks/sql/result_set.py:382
  - #4419356679 at src/databricks/sql/backend/sea/result_set.py:178

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/backend/sea/utils/filters.py | 10 ++++++----
 src/databricks/sql/result_set.py                | 12 ++++++------
 tests/unit/test_filters.py                      |  4 ++--
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/databricks/sql/backend/sea/utils/filters.py b/src/databricks/sql/backend/sea/utils/filters.py
index f125c409a..662c8b243 100644
--- a/src/databricks/sql/backend/sea/utils/filters.py
+++ b/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:
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index bdba7e164..a74c42394 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -373,10 +373,7 @@ def _fetchone_buffered(
         """
         if self._row_buffer_count() == 0:
             cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
-            target = getattr(
-                self, "_fetchone_refill_rows", _FETCHONE_BUFFER_INITIAL_ROWS
-            )
-            size = max(min(target, cap), 1)
+            size = max(min(self._fetchone_refill_rows, cap), 1)
             self._fetchone_refill_rows = min(
                 size * _FETCHONE_BUFFER_GROWTH_FACTOR, _FETCHONE_BUFFER_MAX_ROWS
             )
@@ -396,12 +393,15 @@ def _fetchone_buffered(
                 # Warn once per result set; later failing batches (e.g. a column
                 # type that fails in every row) log at DEBUG to avoid a WARNING
                 # per refill for callers that catch and keep fetching.
-                warned = getattr(self, "_fetchone_conversion_warned", False)
                 fallback = self._convert_rows_individually(
                     raw,
                     convert,
                     batch_error,
-                    log_level=logging.DEBUG if warned else logging.WARNING,
+                    log_level=(
+                        logging.DEBUG
+                        if self._fetchone_conversion_warned
+                        else logging.WARNING
+                    ),
                 )
                 if fallback is None:
                     raise
diff --git a/tests/unit/test_filters.py b/tests/unit/test_filters.py
index 4efe51f3e..c6ada1475 100644
--- a/tests/unit/test_filters.py
+++ b/tests/unit/test_filters.py
@@ -16,9 +16,9 @@ def setUp(self):
         # Create a mock SeaResultSet
         self.mock_sea_result_set = MagicMock()
 
-        # Set up the remaining_rows method on the results attribute
+        # Set up fetchall_json, which the filters use to read all remaining rows
         self.mock_sea_result_set.results = MagicMock()
-        self.mock_sea_result_set.results.remaining_rows.return_value = [
+        self.mock_sea_result_set.fetchall_json.return_value = [
             ["catalog1", "schema1", "table1", "owner1", "2023-01-01", "TABLE", ""],
             ["catalog1", "schema1", "table2", "owner1", "2023-01-01", "VIEW", ""],
             [

From 62fb86efdbfe3df8a146777209c5fe3d8b965a9b Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:50:28 +0000
Subject: [PATCH 23/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4419399913 at src/databricks/sql/result_set.py:395

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index a74c42394..b8ec922fa 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -390,6 +390,15 @@ def _fetchone_buffered(
                 # as before batching, only the row(s) that fail to convert are
                 # lost: each failing row is parked as a _FailedRow and re-raises
                 # its own error when it is reached.
+                # Why not just drop the batch and re-raise (as fetchmany does)?
+                # fetchone() used to convert one row at a time, so callers that
+                # catch a conversion error and keep iterating lost only that row.
+                # Dropping the whole refill would silently lose up to
+                # _FETCHONE_BUFFER_MAX_ROWS good rows per bad one, a behaviour
+                # change hidden behind a perf fix. Bisection keeps the old
+                # semantics; its ~2n-conversion worst case is bounded by the ramp
+                # reset below, and _FailedRow tracebacks are freed as they are
+                # consumed (or with the buffer).
                 # Warn once per result set; later failing batches (e.g. a column
                 # type that fails in every row) log at DEBUG to avoid a WARNING
                 # per refill for callers that catch and keep fetching.

From 313aa5f549ce0ba2c7ec012ee715f574dd23011d Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 00:59:29 +0000
Subject: [PATCH 24/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4419434164 at src/databricks/sql/result_set.py:383
  - #4419434169 at src/databricks/sql/result_set.py:307

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 20 +++++++++++++++-----
 tests/unit/test_fetches.py       | 10 +++++++++-
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index b8ec922fa..e3a518baa 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -302,18 +302,21 @@ def convert_range(lo: int, hi: int, error: Optional[Exception]) -> None:
         convert_range(0, n, batch_error)
         failed = sum(isinstance(r, _FailedRow) for r in rows)
         # One summary line, no traceback: the error itself still reaches the caller
-        # when the failing row is fetched. The full traceback is only logged at DEBUG.
+        # when the failing row is fetched. Only the exception type is logged here,
+        # since conversion error messages often quote the offending result value;
+        # the message and full traceback are only logged at DEBUG.
         # ``log_level`` lets the caller log this at WARNING only once per result set.
         logger.log(
             log_level,
             "fetchone(): failed to convert %d of %d row(s) in a batch; each failing "
-            "row raises when fetched (%s: %s)",
+            "row raises when fetched (%s)",
             failed,
             n,
             type(batch_error).__name__,
-            batch_error,
         )
-        logger.debug("fetchone(): batch conversion error", exc_info=batch_error)
+        logger.debug(
+            "fetchone(): batch conversion error: %s", batch_error, exc_info=batch_error
+        )
         return rows
 
     def _cap_to_current_chunk(self, size: int) -> int:
@@ -369,7 +372,11 @@ def _fetchone_buffered(
         If conversion yields a different number of rows than the raw chunk held,
         ``InternalError`` is raised and that whole chunk is dropped (the queue and
         ``rownumber`` have already moved past it). The result set should be
-        considered unusable after that error.
+        considered unusable after that error. The same holds when conversion
+        raises one of ``_NON_DATA_CONVERSION_ERRORS`` (e.g. ``MemoryError``,
+        ``RecursionError``): the error is re-raised and the whole refill chunk
+        (up to ``_FETCHONE_BUFFER_MAX_ROWS`` rows, still counted in ``rownumber``)
+        is lost, so callers should not keep fetching from this result set.
         """
         if self._row_buffer_count() == 0:
             cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
@@ -383,6 +390,9 @@ def _fetchone_buffered(
             except _NON_DATA_CONVERSION_ERRORS:
                 # Resource / interpreter failures say nothing about individual
                 # rows; retrying smaller slices would only repeat (or worsen) them.
+                # The queue has already moved past ``raw`` and it is not buffered,
+                # so the whole refill chunk is lost: as with the drift error below,
+                # the result set is unusable after this.
                 raise
             except Exception as batch_error:
                 # The queue has already moved past this batch, so it can't be
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 1d977f9e6..e94cb8579 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -738,6 +738,11 @@ def test_fetchone_convert_failure_warns_once_per_result_set(self):
             [r.levelno for r in summaries],
             [logging.WARNING] + [logging.DEBUG] * 8,
         )
+        # The error message may quote result data, so it must stay out of WARNING.
+        for r in logs.records:
+            if r.levelno >= logging.WARNING:
+                self.assertNotIn("bad value", r.getMessage())
+                self.assertIn("ValueError", r.getMessage())
 
     def test_fetchone_all_failing_batch_resets_refill_ramp(self):
         from unittest.mock import patch
@@ -821,7 +826,10 @@ def convert(self_, table):
         self.assertEqual(len(calls), 2 * 20 - 1)
         self.assertEqual(sum(1 for c in calls if c == 1), 20)
         self.assertTrue(any("20 of 20 row(s)" in m for m in logs.output))
-        self.assertTrue(any("batch failure" in m for m in logs.output))
+        # Only the exception type is logged at WARNING; the message (which may quote
+        # result data) is kept out of it.
+        self.assertTrue(any("(ValueError)" in m for m in logs.output))
+        self.assertFalse(any("batch failure" in m for m in logs.output))
         self.assertEqual(rs.fetchone()[0], 20)
 
     def test_fetchone_all_failing_max_batch_conversion_count_is_bounded(self):

From a6d6ee55f48e10655189986af2db68583e72406c Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 01:07:01 +0000
Subject: [PATCH 25/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4419479174 at src/databricks/sql/backend/sea/utils/filters.py:179
  - #4419479176 at src/databricks/sql/backend/kernel/result_set.py:211

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 .../sql/backend/kernel/result_set.py          |  12 +-
 tests/unit/test_filters.py                    | 118 ++++++++++++++++++
 2 files changed, 127 insertions(+), 3 deletions(-)

diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index 04754041d..ddae531e5 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -209,9 +209,15 @@ def _fetch_available(self, size: int) -> pyarrow.Table:
         return self._take_buffered(min(size, self._buffered_rows()))
 
     @staticmethod
-    def _prepend_buffered(
+    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:
@@ -220,7 +226,7 @@ def _prepend_buffered(
 
     def fetchall_arrow(self) -> pyarrow.Table:
         buffered = self._row_buffer_take_raw()
-        return self._prepend_buffered(buffered, self._drain())
+        return self._concat_buffered_arrow(buffered, self._drain())
 
     def fetchmany_arrow(self, size: int) -> pyarrow.Table:
         if size < 0:
@@ -229,7 +235,7 @@ def fetchmany_arrow(self, size: int) -> pyarrow.Table:
         remaining = size - (buffered.num_rows if buffered is not None else 0)
         if buffered is not None and remaining == 0:
             return buffered
-        return self._prepend_buffered(
+        return self._concat_buffered_arrow(
             buffered, self._fetchmany_arrow_unbuffered(remaining)
         )
 
diff --git a/tests/unit/test_filters.py b/tests/unit/test_filters.py
index c6ada1475..7036044c0 100644
--- a/tests/unit/test_filters.py
+++ b/tests/unit/test_filters.py
@@ -5,6 +5,8 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
+import pytest
+
 from databricks.sql.backend.sea.utils.filters import ResultSetFilter
 
 
@@ -153,5 +155,121 @@ def test_filter_tables_by_type(self):
             )
 
 
+class TestResultSetFilterWithFetchoneBuffer(unittest.TestCase):
+    """Rows buffered by fetchone() on a real SeaResultSet must reach the filter."""
+
+    N_ROWS = 12
+
+    def _row_type(self, i):
+        return "TABLE" if i % 2 == 0 else "VIEW"
+
+    def _make_result_set(self, queue, result_format):
+        from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
+        from databricks.sql.backend.sea.result_set import SeaResultSet
+        from databricks.sql.backend.types import CommandId, CommandState
+
+        connection = MagicMock()
+        connection.open = True
+        connection.disable_pandas = True
+        sea_client = MagicMock()
+        sea_client.max_download_threads = 10
+
+        execute_response = MagicMock()
+        execute_response.command_id = CommandId.from_sea_statement_id("stmt-1")
+        execute_response.status = CommandState.SUCCEEDED
+        execute_response.has_been_closed_server_side = False
+        execute_response.description = [
+            ("table_name", "string", None, None, None, None, None),
+            ("table_type", "string", None, None, None, None, None),
+        ]
+        execute_response.is_staging_operation = False
+        execute_response.lz4_compressed = False
+        execute_response.arrow_schema_bytes = None
+
+        manifest = ResultManifest(
+            format=result_format.value,
+            schema={},
+            total_row_count=self.N_ROWS,
+            total_byte_count=-1,
+            total_chunk_count=-1,
+        )
+        with patch(
+            "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue",
+            return_value=queue,
+        ):
+            return SeaResultSet(
+                connection=connection,
+                execute_response=execute_response,
+                sea_client=sea_client,
+                result_data=ResultData(data=None, external_links=None),
+                manifest=manifest,
+                buffer_size_bytes=1000,
+                arraysize=100,
+            )
+
+    def _consume_two_rows(self, rs):
+        # 1st fetchone() refills 1 row, 2nd refills 10: rows 2..10 stay buffered
+        # and row 11 is still in the queue.
+        self.assertEqual(rs.fetchone()[0], "t0")
+        self.assertEqual(rs.fetchone()[0], "t1")
+        self.assertGreater(rs._row_buffer_count(), 0)
+
+    def _expected_names(self):
+        return [
+            "t{}".format(i)
+            for i in range(2, self.N_ROWS)
+            if self._row_type(i) == "TABLE"
+        ]
+
+    def test_json_filter_includes_rows_buffered_by_fetchone(self):
+        from databricks.sql.backend.sea.queue import JsonQueue
+        from databricks.sql.backend.sea.utils.constants import ResultFormat
+
+        data = [["t{}".format(i), self._row_type(i)] for i in range(self.N_ROWS)]
+        rs = self._make_result_set(JsonQueue(data), ResultFormat.JSON_ARRAY)
+        self._consume_two_rows(rs)
+
+        with patch(
+            "databricks.sql.backend.sea.result_set.SeaResultSet"
+        ) as mock_sea_result_set_class:
+            ResultSetFilter._filter_json_result_set(
+                rs, 1, ["TABLE"], case_sensitive=True
+            )
+            result_data = mock_sea_result_set_class.call_args.kwargs["result_data"]
+
+        self.assertEqual([r[0] for r in result_data.data], self._expected_names())
+        self.assertEqual(rs.rownumber, self.N_ROWS)
+
+    def test_arrow_filter_includes_rows_buffered_by_fetchone(self):
+        pyarrow = pytest.importorskip("pyarrow")
+        from databricks.sql.backend.sea.utils.constants import ResultFormat
+        from databricks.sql.utils import ArrowQueue
+
+        table = pyarrow.Table.from_pydict(
+            {
+                "table_name": ["t{}".format(i) for i in range(self.N_ROWS)],
+                "table_type": [self._row_type(i) for i in range(self.N_ROWS)],
+            }
+        )
+        rs = self._make_result_set(
+            ArrowQueue(table, self.N_ROWS), ResultFormat.ARROW_STREAM
+        )
+        self._consume_two_rows(rs)
+
+        with patch(
+            "databricks.sql.backend.sea.result_set.SeaResultSet"
+        ) as mock_sea_result_set_class:
+            ResultSetFilter._filter_arrow_result_set(
+                rs, 1, ["TABLE"], case_sensitive=True
+            )
+            result_data = mock_sea_result_set_class.call_args.kwargs["result_data"]
+
+        filtered = pyarrow.ipc.open_stream(result_data.attachment).read_all()
+        self.assertEqual(
+            filtered.column("table_name").to_pylist(), self._expected_names()
+        )
+        self.assertEqual(rs.rownumber, self.N_ROWS)
+
+
 if __name__ == "__main__":
     unittest.main()

From 37a9c535546f9d8dc5ab1e05812d195d81f8a758 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 01:16:09 +0000
Subject: [PATCH 26/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4419518440 at src/databricks/sql/result_set.py:417
  - #4419518444 at src/databricks/sql/result_set.py:258

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 14 ++++++++--
 tests/unit/test_fetches.py       | 47 ++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index e3a518baa..74aacadc4 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -381,10 +381,12 @@ def _fetchone_buffered(
         if self._row_buffer_count() == 0:
             cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
             size = max(min(self._fetchone_refill_rows, cap), 1)
+            raw = fetch_raw(size)
+            # Grow the ramp only once the fetch succeeded, so a transient fetch
+            # error (and a retried fetchone()) does not skip a ramp step.
             self._fetchone_refill_rows = min(
                 size * _FETCHONE_BUFFER_GROWTH_FACTOR, _FETCHONE_BUFFER_MAX_ROWS
             )
-            raw = fetch_raw(size)
             try:
                 rows = convert(raw)
             except _NON_DATA_CONVERSION_ERRORS:
@@ -436,7 +438,15 @@ def _fetchone_buffered(
             # that yields [] for a non-empty chunk raises instead of silently
             # ending iteration and losing those rows.
             raw_count = self._raw_row_count(raw)
-            if raw_count is not None and raw_count != len(rows):
+            if raw_count is None:
+                # Every real queue type (Arrow Table, ColumnTable, list) has a row
+                # count, so an unknown chunk type is itself a bug. Fail loudly
+                # rather than skip the drift check and risk rownumber misalignment.
+                raise InternalError(
+                    "fetchone(): cannot determine the row count of a raw chunk of "
+                    "type {}".format(type(raw).__name__)
+                )
+            if raw_count != len(rows):
                 # fetch_raw advanced _next_row_index (and the Thrift fetch offset)
                 # by raw_count; the buffer must line up 1:1 with the raw chunk or
                 # rownumber and fetchmany_arrow slicing would silently drift.
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index e94cb8579..d91277f2e 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -869,6 +869,30 @@ def convert(self_, table):
         self.assertLessEqual(sum(calls), n * (n - 1).bit_length() + n)
         self.assertEqual(rs._row_buffer_count(), n - 1)
 
+    def test_fetchone_failed_fetch_does_not_advance_refill_ramp(self):
+        from unittest.mock import patch
+
+        rs = self.make_columnar_result_set(100, arraysize=100, warm=False)
+        initial = rs._fetchone_refill_rows
+        real_next_n_rows = rs.results.next_n_rows
+        with patch.object(
+            rs.results,
+            "next_n_rows",
+            side_effect=[RuntimeError("transient"), real_next_n_rows],
+        ):
+            with self.assertRaises(RuntimeError):
+                rs.fetchone()
+        # The failed fetch must not grow the ramp, so a retry refills the same size.
+        self.assertEqual(rs._fetchone_refill_rows, initial)
+        sizes = []
+        with patch.object(
+            rs.results,
+            "next_n_rows",
+            side_effect=lambda n: sizes.append(n) or real_next_n_rows(n),
+        ):
+            self.assertEqual(rs.fetchone()[0], 0)
+        self.assertEqual(sizes, [initial])
+
     def test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows(self):
         from unittest.mock import patch
 
@@ -962,6 +986,29 @@ def test_fetchone_empty_conversion_of_nonempty_chunk_fails_loudly(self):
             with self.assertRaises(InternalError):
                 rs.fetchone()
 
+    def test_fetchone_unknown_raw_chunk_type_fails_loudly(self):
+        from unittest.mock import patch
+        from databricks.sql.exc import InternalError
+
+        class _NoRowCount:
+            pass
+
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_fetch_available_from_queue",
+            autospec=True,
+            return_value=_NoRowCount(),
+        ), patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            return_value=[("x",)],
+        ):
+            with self.assertRaises(InternalError) as ctx:
+                rs.fetchone()
+        self.assertIn("_NoRowCount", str(ctx.exception))
+
     @pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
     def test_fetchone_then_fetchall_arrow_keeps_order(self):
         rs = self.make_columnar_result_set(10, arraysize=4)

From c0a02a017b8f5aab62dc8c49db1c4a2917a3da0f Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 01:28:29 +0000
Subject: [PATCH 27/29] ai: apply changes for #1007 (2 review threads)

Addresses:
  - #4419566436 at src/databricks/sql/result_set.py:391
  - #4419566437 at src/databricks/sql/result_set.py:318

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py | 48 +++++++++++++++++++++++++++++---
 tests/unit/test_fetches.py       | 35 +++++++++++++++++++++++
 2 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index 74aacadc4..edcd23083 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -28,6 +28,14 @@
 
 logger = logging.getLogger(__name__)
 
+# Records that may contain result data go to the non-propagating "unsafe" logger
+# (configured the same way in thrift_backend.py), so enabling DEBUG on
+# ``databricks.sql`` does not leak row values into logs.
+unsafe_logger = logging.getLogger("databricks.sql.unsafe")
+unsafe_logger.propagate = False
+if not unsafe_logger.handlers:
+    unsafe_logger.addHandler(logging.NullHandler())
+
 # Upper bound on how many rows a fetchone() refill pulls and converts at once. The
 # refill uses ``arraysize`` but is capped here so a cursor with a large arraysize
 # (default 100000) doesn't materialise a full page just to return one row.
@@ -130,6 +138,10 @@ def __init__(
         # Size of the next fetchone() refill; grows on each refill (see
         # ``_FETCHONE_BUFFER_GROWTH_FACTOR``).
         self._fetchone_refill_rows = _FETCHONE_BUFFER_INITIAL_ROWS
+        # Set when a fetchone() refill chunk was read from the queue but lost
+        # before it could be buffered (e.g. KeyboardInterrupt / MemoryError during
+        # conversion). Later fetches raise instead of silently skipping those rows.
+        self._lost_rows_error: Optional[str] = None
 
     def __iter__(self):
         while True:
@@ -183,6 +195,11 @@ def rownumber(self):
 
     # ----- fetchone() row buffer helpers -----
 
+    def _check_no_lost_rows(self) -> None:
+        """Raise if an earlier fetchone() refill lost rows (see ``_lost_rows_error``)."""
+        if self._lost_rows_error is not None:
+            raise InternalError(self._lost_rows_error)
+
     def _row_buffer_count(self) -> int:
         return len(self._row_buffer_rows) - self._row_buffer_pos
 
@@ -211,6 +228,7 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
         all of its rows or raises), so ``rownumber`` stays aligned with what is
         left in the buffer.
         """
+        self._check_no_lost_rows()
         if self._row_buffer_count() == 0:
             return []
         rows = self._row_buffer_rows
@@ -304,7 +322,7 @@ def convert_range(lo: int, hi: int, error: Optional[Exception]) -> None:
         # One summary line, no traceback: the error itself still reaches the caller
         # when the failing row is fetched. Only the exception type is logged here,
         # since conversion error messages often quote the offending result value;
-        # the message and full traceback are only logged at DEBUG.
+        # the message and full traceback only go to the non-propagating unsafe logger.
         # ``log_level`` lets the caller log this at WARNING only once per result set.
         logger.log(
             log_level,
@@ -314,7 +332,7 @@ def convert_range(lo: int, hi: int, error: Optional[Exception]) -> None:
             n,
             type(batch_error).__name__,
         )
-        logger.debug(
+        unsafe_logger.debug(
             "fetchone(): batch conversion error: %s", batch_error, exc_info=batch_error
         )
         return rows
@@ -340,6 +358,7 @@ def _row_buffer_take_raw(self, n: Optional[int] = None) -> Any:
         Returns None when the buffer is empty, otherwise a slice of the raw chunk
         (same type as what the queue returned: Arrow table, ColumnTable or list).
         """
+        self._check_no_lost_rows()
         if self._row_buffer_count() == 0:
             return None
         raw = self._row_buffer_raw
@@ -377,7 +396,12 @@ def _fetchone_buffered(
         ``RecursionError``): the error is re-raised and the whole refill chunk
         (up to ``_FETCHONE_BUFFER_MAX_ROWS`` rows, still counted in ``rownumber``)
         is lost, so callers should not keep fetching from this result set.
+        The same applies to any ``BaseException`` (e.g. ``KeyboardInterrupt``)
+        raised during conversion. In all these cases every later fetch on this
+        result set raises ``InternalError`` rather than silently skipping the
+        lost rows.
         """
+        self._check_no_lost_rows()
         if self._row_buffer_count() == 0:
             cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
             size = max(min(self._fetchone_refill_rows, cap), 1)
@@ -387,6 +411,20 @@ def _fetchone_buffered(
             self._fetchone_refill_rows = min(
                 size * _FETCHONE_BUFFER_GROWTH_FACTOR, _FETCHONE_BUFFER_MAX_ROWS
             )
+            # The queue (and rownumber) has already moved past ``raw``, but it is
+            # only buffered once conversion succeeds. Mark the result set as having
+            # lost rows until then, so any exception escaping below (including
+            # KeyboardInterrupt) makes later fetches fail loudly instead of
+            # resuming from the next chunk with a gap.
+            pending = self._raw_row_count(raw)
+            if pending != 0:
+                self._lost_rows_error = (
+                    "{} row(s) were lost because an earlier fetchone() failed while "
+                    "converting them; this result set can no longer be read "
+                    "consistently, re-execute the query".format(
+                        "Some" if pending is None else pending
+                    )
+                )
             try:
                 rows = convert(raw)
             except _NON_DATA_CONVERSION_ERRORS:
@@ -451,8 +489,8 @@ def _fetchone_buffered(
                 # by raw_count; the buffer must line up 1:1 with the raw chunk or
                 # rownumber and fetchmany_arrow slicing would silently drift.
                 # The raw chunk is not buffered, so its raw_count rows are lost
-                # (yet counted in rownumber): treat the result set as unusable
-                # after this error; later fetches resume from the next chunk.
+                # (yet counted in rownumber): the result set is unusable after
+                # this error and later fetches raise (see ``_lost_rows_error``).
                 raise InternalError(
                     "fetchone(): converted {} row(s) from a chunk of {} raw row(s)".format(
                         len(rows), raw_count
@@ -464,6 +502,8 @@ def _fetchone_buffered(
             self._row_buffer_rows = rows
             self._row_buffer_pos = 0
             self._next_row_index -= len(rows)
+            # ``raw`` is buffered now, so no rows were lost.
+            self._lost_rows_error = None
         return self._row_buffer_take_rows(1)[0]
 
     @property
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index d91277f2e..80b7b7eb8 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -12,6 +12,7 @@
 from databricks.sql.utils import ArrowQueue
 from databricks.sql.backend.thrift_backend import ThriftDatabricksClient
 from databricks.sql.result_set import ThriftResultSet
+from databricks.sql.exc import InternalError
 
 
 class _StubArrowQueue:
@@ -789,6 +790,40 @@ def convert(self_, table):
                 rs.fetchone()
         self.assertEqual(calls, [4])
         self.assertEqual(rs._row_buffer_count(), 0)
+        # The lost refill chunk must not be skipped silently on later fetches.
+        with self.assertRaises(InternalError):
+            rs.fetchone()
+
+    def test_fetchone_interrupted_conversion_poisons_later_fetches(self):
+        from unittest.mock import patch
+
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            side_effect=KeyboardInterrupt(),
+        ):
+            with self.assertRaises(KeyboardInterrupt):
+                rs.fetchone()
+        # Conversion works again, but the interrupted chunk is gone: every fetch
+        # must raise instead of resuming from the next chunk with a gap.
+        for fetch in (
+            rs.fetchone,
+            lambda: rs.fetchmany(2),
+            rs.fetchall,
+            lambda: rs.fetchmany_arrow(2),
+            rs.fetchall_arrow,
+        ):
+            with self.assertRaises(InternalError) as ctx:
+                fetch()
+            self.assertIn("4 row(s) were lost", str(ctx.exception))
+
+    def test_fetchone_successful_refill_does_not_poison(self):
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        rows = [rs.fetchone() for _ in range(3)]
+        self.assertEqual([r[0] for r in rows], [0, 1, 2])
+        self.assertIsNone(rs._lost_rows_error)
+        self.assertEqual([r[0] for r in rs.fetchall()], list(range(3, 10)))
 
     def test_fetchone_convert_failure_on_every_row_bisects_to_single_rows(self):
         from unittest.mock import patch

From 4a4bc039044a32587d1c53f362277711c02de636 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 01:37:50 +0000
Subject: [PATCH 28/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4419629658 at src/databricks/sql/result_set.py:312

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/result_set.py |  4 ++++
 tests/unit/test_fetches.py       | 28 ++++++++++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index edcd23083..fbea1947f 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -308,6 +308,10 @@ def convert_range(lo: int, hi: int, error: Optional[Exception]) -> None:
                 try:
                     rows.extend(convert(piece))
                     return
+                except _NON_DATA_CONVERSION_ERRORS:
+                    # Resource / interpreter failures are not row-specific: stop
+                    # bisecting and propagate, as the top-level refill does.
+                    raise
                 except Exception as e:
                     error = e
             if hi - lo == 1:
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index 80b7b7eb8..cd3e892c7 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -794,6 +794,34 @@ def convert(self_, table):
         with self.assertRaises(InternalError):
             rs.fetchone()
 
+    def test_fetchone_memory_error_during_bisection_stops_bisecting(self):
+        from unittest.mock import patch
+
+        calls = []
+
+        def convert(self_, table):
+            calls.append(table.num_rows)
+            if len(calls) == 1:
+                raise ValueError("bad value")
+            raise MemoryError()
+
+        rs = self.make_columnar_result_set(10, arraysize=4)
+        with patch.object(
+            ThriftResultSet,
+            "_convert_columnar_table",
+            autospec=True,
+            side_effect=convert,
+        ):
+            with self.assertRaises(MemoryError):
+                rs.fetchone()
+        # The full batch fails with a data error, the first half hits MemoryError:
+        # bisection stops there instead of converting more slices.
+        self.assertEqual(calls, [4, 2])
+        self.assertEqual(rs._row_buffer_count(), 0)
+        # The lost refill chunk must not be skipped silently on later fetches.
+        with self.assertRaises(InternalError):
+            rs.fetchone()
+
     def test_fetchone_interrupted_conversion_poisons_later_fetches(self):
         from unittest.mock import patch
 

From 02b8baf39a9cbb31dafd0e2445c0e3f09905ed48 Mon Sep 17 00:00:00 2001
From: "peco-engineer-bot[bot]"
 <peco-engineer-bot[bot]@users.noreply.github.com>
Date: Fri, 30 Oct 2026 01:45:59 +0000
Subject: [PATCH 29/29] ai: apply changes for #1007 (1 review thread)

Addresses:
  - #4419671350 at src/databricks/sql/result_set.py:567

Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
---
 src/databricks/sql/backend/kernel/result_set.py |  5 +++--
 src/databricks/sql/result_set.py                |  4 +++-
 tests/unit/test_fetches.py                      | 17 +++++++++++++++++
 tests/unit/test_kernel_result_set.py            | 10 ++++++++++
 4 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py
index ddae531e5..e5a32a37c 100644
--- a/src/databricks/sql/backend/kernel/result_set.py
+++ b/src/databricks/sql/backend/kernel/result_set.py
@@ -273,9 +273,10 @@ def close(self) -> None:
         connection-close would otherwise issue a kernel call into an
         already-disposed session.
         """
-        # Drop rows parked by fetchone() so post-close fetches return
-        # nothing, matching the cleared batch FIFO below.
+        # 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:
diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py
index fbea1947f..bf1761230 100644
--- a/src/databricks/sql/result_set.py
+++ b/src/databricks/sql/result_set.py
@@ -563,8 +563,10 @@ def close(self) -> None:
             if isinstance(e.args[1], CursorAlreadyClosedError):
                 logger.info("Operation was canceled by a prior request")
         finally:
-            # Drop rows parked by fetchone() so post-close fetches return nothing.
+            # Drop rows parked by fetchone() (and any lost-rows poison) so
+            # post-close fetches return nothing.
             self._row_buffer_clear()
+            self._lost_rows_error = None
             self.has_been_closed_server_side = True
             self.status = CommandState.CLOSED
 
diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py
index cd3e892c7..bab2c2aa0 100644
--- a/tests/unit/test_fetches.py
+++ b/tests/unit/test_fetches.py
@@ -390,6 +390,23 @@ def test_close_drops_fetchone_row_buffer(self):
             self.assertEqual(fetch_after_close(rs), [])
             self.assertIsNone(rs.fetchone())
 
+    def test_close_clears_lost_rows_error(self):
+        # A result set poisoned by a lost fetchone() refill must behave like any
+        # other closed result set: post-close fetches return nothing, not raise.
+        from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
+
+        rs = self.make_dummy_result_set_from_initial_results([[1], [2], [3]])
+        # Warm refill pulls every row into the buffer, so the queue is empty.
+        rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
+        self.assertSequenceEqual(rs.fetchone(), [1])
+        rs._lost_rows_error = "lost"
+        with self.assertRaises(InternalError):
+            rs.fetchone()
+        rs.close()
+        self.assertIsNone(rs._lost_rows_error)
+        self.assertIsNone(rs.fetchone())
+        self.assertEqual(rs.fetchall(), [])
+
     def test_fetchone_refill_does_not_fetch_extra_pages(self):
         # A refill must not make extra round-trips to fill ``arraysize`` when
         # server pages are smaller (e.g. byte-limited): it serves what the current
diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py
index f9572b018..ef2dcd6d7 100644
--- a/tests/unit/test_kernel_result_set.py
+++ b/tests/unit/test_kernel_result_set.py
@@ -288,6 +288,16 @@ def test_close_drops_fetchone_row_buffer(int_schema):
     assert rs.fetchall() == []
 
 
+def test_close_clears_lost_rows_error(int_schema):
+    # A poisoned result set must return nothing after close(), not keep raising.
+    rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
+    rs._lost_rows_error = "lost"
+    rs.close()
+    assert rs._lost_rows_error is None
+    assert rs.fetchone() is None
+    assert rs.fetchall() == []
+
+
 def test_empty_stream(int_schema):
     rs = _make_rs(_FakeKernelHandle(int_schema, []))
     assert rs.fetchone() is None