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>

Author: peco-engineer-bot[bot]

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

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)