Skip to content

Commit 4601240

Browse files
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>
1 parent 77e1090 commit 4601240

4 files changed

Lines changed: 89 additions & 35 deletions

File tree

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

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -245,17 +245,15 @@ def fetchone(self) -> Optional[Row]:
245245
def fetchmany(self, size: int) -> List[Row]:
246246
if size < 0:
247247
raise ValueError(f"fetchmany size must be >= 0, got {size}")
248-
buffered = self._row_buffer_take_rows(size)
249-
remaining = size - len(buffered)
250-
if remaining == 0:
251-
return buffered
252-
rest = self._convert_arrow_table(self._fetchmany_arrow_unbuffered(remaining))
253-
return buffered + rest if buffered else rest
248+
return self._fetch_rows_with_buffer(
249+
size,
250+
lambda n: self._convert_arrow_table(self._fetchmany_arrow_unbuffered(n)),
251+
)
254252

255253
def fetchall(self) -> List[Row]:
256-
buffered = self._row_buffer_take_rows()
257-
rest = self._convert_arrow_table(self._drain())
258-
return buffered + rest if buffered else rest
254+
return self._fetch_rows_with_buffer(
255+
None, lambda _: self._convert_arrow_table(self._drain())
256+
)
259257

260258
def close(self) -> None:
261259
"""Close the underlying kernel handle and notify the backend.

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

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -279,15 +279,13 @@ def fetchmany(self, size: int) -> List[Row]:
279279

280280
if size < 0:
281281
raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
282-
buffered = self._row_buffer_take_rows(size)
283-
remaining = size - len(buffered)
284-
if buffered and remaining == 0:
285-
return buffered
286282
if isinstance(self.results, JsonQueue):
287-
rest = self._create_json_table(self.fetchmany_json(remaining))
288-
else:
289-
rest = self._convert_arrow_table(self.fetchmany_arrow(remaining))
290-
return buffered + rest if buffered else rest
283+
return self._fetch_rows_with_buffer(
284+
size, lambda n: self._create_json_table(self.fetchmany_json(n))
285+
)
286+
return self._fetch_rows_with_buffer(
287+
size, lambda n: self._convert_arrow_table(self.fetchmany_arrow(n))
288+
)
291289

292290
def fetchall(self) -> List[Row]:
293291
"""
@@ -297,9 +295,10 @@ def fetchall(self) -> List[Row]:
297295
List of Row objects containing all remaining rows
298296
"""
299297

300-
buffered = self._row_buffer_take_rows()
301298
if isinstance(self.results, JsonQueue):
302-
rest = self._create_json_table(self.fetchall_json())
303-
else:
304-
rest = self._convert_arrow_table(self.fetchall_arrow())
305-
return buffered + rest if buffered else rest
299+
return self._fetch_rows_with_buffer(
300+
None, lambda _: self._create_json_table(self.fetchall_json())
301+
)
302+
return self._fetch_rows_with_buffer(
303+
None, lambda _: self._convert_arrow_table(self.fetchall_arrow())
304+
)

‎src/databricks/sql/result_set.py‎

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
_FETCHONE_BUFFER_INITIAL_ROWS = 1
4242
_FETCHONE_BUFFER_GROWTH_FACTOR = 10
4343

44+
4445
class _FailedRow:
4546
"""Placeholder in the fetchone() row buffer for a row that failed to convert."""
4647

@@ -226,6 +227,26 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
226227
raise failed.error.with_traceback(failed.traceback)
227228
return rows[start : start + length]
228229

230+
def _fetch_rows_with_buffer(
231+
self,
232+
size: Optional[int],
233+
fetch_rest: Callable[[Optional[int]], List[Row]],
234+
) -> List[Row]:
235+
"""Return ``size`` rows (all if None), draining the fetchone() buffer first.
236+
237+
Every row-level fetchmany()/fetchall() goes through here, so rows the
238+
fetchone() buffer already holds are always returned before newly fetched
239+
ones. ``fetch_rest(n)`` must fetch up to ``n`` (all if None) rows
240+
while bypassing the buffer. It is not called when the buffer already
241+
covers ``size``.
242+
"""
243+
buffered = self._row_buffer_take_rows(size)
244+
remaining = None if size is None else size - len(buffered)
245+
if buffered and remaining == 0:
246+
return buffered
247+
rest = fetch_rest(remaining)
248+
return buffered + rest if buffered else rest
249+
229250
@staticmethod
230251
def _raw_row_count(raw: Any) -> Optional[int]:
231252
"""Row count of a raw chunk (Arrow table, ColumnTable or list); None if unknown."""
@@ -774,12 +795,13 @@ def fetchall(self) -> List[Row]:
774795
"""
775796
Fetch all (remaining) rows of a query result, returning them as a list of rows.
776797
"""
777-
buffered = self._row_buffer_take_rows()
778798
if isinstance(self.results, ColumnQueue):
779-
rest = self._convert_columnar_table(self.fetchall_columnar())
780-
else:
781-
rest = self._convert_arrow_table(self.fetchall_arrow())
782-
return buffered + rest if buffered else rest
799+
return self._fetch_rows_with_buffer(
800+
None, lambda _: self._convert_columnar_table(self.fetchall_columnar())
801+
)
802+
return self._fetch_rows_with_buffer(
803+
None, lambda _: self._convert_arrow_table(self.fetchall_arrow())
804+
)
783805

784806
def fetchmany(self, size: int) -> List[Row]:
785807
"""
@@ -789,15 +811,13 @@ def fetchmany(self, size: int) -> List[Row]:
789811
"""
790812
if size < 0:
791813
raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
792-
buffered = self._row_buffer_take_rows(size)
793-
remaining = size - len(buffered)
794-
if buffered and remaining == 0:
795-
return buffered
796814
if isinstance(self.results, ColumnQueue):
797-
rest = self._convert_columnar_table(self.fetchmany_columnar(remaining))
798-
else:
799-
rest = self._convert_arrow_table(self.fetchmany_arrow(remaining))
800-
return buffered + rest if buffered else rest
815+
return self._fetch_rows_with_buffer(
816+
size, lambda n: self._convert_columnar_table(self.fetchmany_columnar(n))
817+
)
818+
return self._fetch_rows_with_buffer(
819+
size, lambda n: self._convert_arrow_table(self.fetchmany_arrow(n))
820+
)
801821

802822
@staticmethod
803823
def _get_schema_description(table_schema_message):

‎tests/unit/test_fetches.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -777,6 +777,43 @@ def convert(self_, table):
777777
self.assertTrue(any("batch failure" in m for m in logs.output))
778778
self.assertEqual(rs.fetchone()[0], 20)
779779

780+
def test_fetchone_all_failing_max_batch_conversion_count_is_bounded(self):
781+
"""Document the worst-case cost of bisecting a full all-failing refill.
782+
783+
When every row fails (a systemic converter / type problem), bisection tries
784+
every node of the tree: exactly ``2n - 1`` conversions for an ``n``-row
785+
refill, i.e. about 2x the one-conversion-per-row cost of fetchone() before
786+
batching. Short-circuiting on "same exception type in both halves" is
787+
deliberately not done, since it would drop good rows that share the error
788+
type (see test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows).
789+
"""
790+
from unittest.mock import patch
791+
from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
792+
793+
n = _FETCHONE_BUFFER_MAX_ROWS
794+
calls = []
795+
796+
def convert(self_, table):
797+
calls.append(table.num_rows)
798+
raise ValueError("systemic failure")
799+
800+
rs = self.make_columnar_result_set(n, arraysize=n)
801+
# Skip the adaptive refill ramp so the first refill is a full batch.
802+
rs._fetchone_refill_rows = n
803+
with patch.object(
804+
ThriftResultSet,
805+
"_convert_columnar_table",
806+
autospec=True,
807+
side_effect=convert,
808+
), self.assertLogs("databricks.sql.result_set", level="WARNING"):
809+
with self.assertRaises(ValueError):
810+
rs.fetchone()
811+
self.assertEqual(calls[0], n)
812+
self.assertEqual(len(calls), 2 * n - 1)
813+
# Total rows converted across all attempts is O(n log n), not O(n^2).
814+
self.assertLessEqual(sum(calls), n * (n - 1).bit_length() + n)
815+
self.assertEqual(rs._row_buffer_count(), n - 1)
816+
780817
def test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows(self):
781818
from unittest.mock import patch
782819

0 commit comments

Comments
 (0)