Skip to content

Commit 05f9f56

Browse files
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>
1 parent 1a44df9 commit 05f9f56

3 files changed

Lines changed: 76 additions & 7 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ def _fetch_available_from_queue(self, size: int) -> Any:
193193
``size`` rows, so a refill that starts near the end of the current chunk
194194
would block on the next download before returning a row that is already
195195
available. Cap the request at what the current chunk still holds (matching
196-
Thrift's ``_fetch_available_from_queue`` and the kernel's ``_fetch_available``).
196+
Thrift's ``_cap_to_current_chunk`` and the kernel's ``_fetch_available``).
197197
Only called while the fetchone() row buffer is empty.
198198
"""
199199
if isinstance(self.results, CloudFetchQueue) and self.results.table:

‎src/databricks/sql/result_set.py‎

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,9 @@
1616
from databricks.sql.client import Connection
1717
from databricks.sql.backend.databricks_client import DatabricksClient
1818
from databricks.sql.types import Row
19-
from databricks.sql.exc import RequestError, CursorAlreadyClosedError
19+
from databricks.sql.exc import InternalError, RequestError, CursorAlreadyClosedError
2020
from databricks.sql.utils import (
21+
CloudFetchQueue,
2122
ColumnTable,
2223
ColumnQueue,
2324
concat_table_chunks,
@@ -196,6 +197,14 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
196197
raise row.error
197198
return taken
198199

200+
@staticmethod
201+
def _raw_row_count(raw: Any) -> Optional[int]:
202+
"""Row count of a raw chunk (Arrow table, ColumnTable or list); None if unknown."""
203+
if raw is None:
204+
return 0
205+
n = len(raw) if isinstance(raw, list) else getattr(raw, "num_rows", None)
206+
return n if isinstance(n, int) else None
207+
199208
@staticmethod
200209
def _convert_rows_individually(
201210
raw: Any, convert: Callable[[Any], List[Row]]
@@ -206,9 +215,7 @@ def _convert_rows_individually(
206215
raised only when that row is handed out. Returns None if ``raw`` has no rows
207216
to split, in which case the caller re-raises the batch error.
208217
"""
209-
if raw is None:
210-
return None
211-
n = len(raw) if isinstance(raw, list) else getattr(raw, "num_rows", 0)
218+
n = ResultSet._raw_row_count(raw)
212219
if not n:
213220
return None
214221
rows: List[Any] = []
@@ -275,6 +282,16 @@ def _fetchone_buffered(
275282
rows = fallback
276283
if not rows:
277284
return None
285+
raw_count = self._raw_row_count(raw)
286+
if raw_count is not None and raw_count != len(rows):
287+
# fetch_raw advanced _next_row_index (and the Thrift fetch offset)
288+
# by raw_count; the buffer must line up 1:1 with the raw chunk or
289+
# rownumber and fetchmany_arrow slicing would silently drift.
290+
raise InternalError(
291+
"fetchone(): converted {} row(s) from a chunk of {} raw row(s)".format(
292+
len(rows), raw_count
293+
)
294+
)
278295
self._row_buffer_raw = raw
279296
self._row_buffer_rows = rows
280297
self._row_buffer_pos = 0
@@ -552,19 +569,31 @@ def _fetch_available_from_queue(self, size: int):
552569
until ``size`` rows are collected: it only refills the queue while it is
553570
empty, so fetchone() makes no more round-trips than it needs to return one
554571
row (server pages can be smaller than ``size`` when they hit the byte limit).
572+
On CloudFetch, ``next_n_rows`` keeps walking downloaded files until it has
573+
``size`` rows, so the request is also capped at what the current chunk still
574+
holds; otherwise a refill near a chunk boundary would block on the next
575+
download before returning a row that is already available.
555576
Only called while the fetchone() row buffer is empty.
556577
"""
557-
results = self.results.next_n_rows(size)
578+
results = self.results.next_n_rows(self._cap_to_current_chunk(size))
558579
while (
559580
results.num_rows == 0
560581
and not self.has_been_closed_server_side
561582
and self.has_more_rows
562583
):
563584
self._fill_results_buffer()
564-
results = self.results.next_n_rows(size)
585+
results = self.results.next_n_rows(self._cap_to_current_chunk(size))
565586
self._next_row_index += results.num_rows
566587
return results
567588

589+
def _cap_to_current_chunk(self, size: int) -> int:
590+
"""Cap ``size`` at the rows left in the current CloudFetch chunk, if any."""
591+
if isinstance(self.results, CloudFetchQueue) and self.results.table:
592+
available = self.results.table.num_rows - self.results.table_row_index
593+
if available > 0:
594+
return min(size, available)
595+
return size
596+
568597
def fetchall_arrow(self) -> "pyarrow.Table":
569598
"""Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
570599
buffered = self._row_buffer_take_raw()

‎tests/unit/test_fetches.py‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -392,6 +392,27 @@ def test_fetchone_refill_does_not_fetch_extra_pages(self):
392392
self.assertIsNone(rs.fetchone())
393393
self.assertEqual(rs.rownumber, 5)
394394

395+
def test_cloudfetch_fetchone_refill_capped_to_current_chunk(self):
396+
# fetchone() must not block on the next CloudFetch file while the current
397+
# chunk still holds rows: the refill is capped at what the chunk has left.
398+
from databricks.sql.utils import ThriftCloudFetchQueue
399+
400+
_, chunk = self.make_arrow_table([[i] for i in range(10)])
401+
queue = Mock(spec=ThriftCloudFetchQueue)
402+
queue.table = chunk
403+
queue.table_row_index = 8
404+
queue.next_n_rows.side_effect = lambda n: chunk.slice(8, n)
405+
rs = self.make_dummy_result_set_from_queue_list(
406+
[queue],
407+
description=[("col0", "integer", None, None, None, None, None)],
408+
)
409+
rs.arraysize = 100
410+
411+
self.assertSequenceEqual(rs.fetchone(), [8])
412+
queue.next_n_rows.assert_called_once_with(2)
413+
self.assertSequenceEqual(rs.fetchone(), [9])
414+
self.assertEqual(rs.rownumber, 2)
415+
395416
# Regression tests for fetchmany_arrow / fetchall_arrow handling of
396417
# the schemaless CloudFetch placeholder.
397418
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):
579600
self.assertEqual(rs._row_buffer_count(), 0)
580601
self.assertEqual(rs.fetchone()[0], 4)
581602

603+
def test_fetchone_raw_converted_count_mismatch_fails_loudly(self):
604+
from unittest.mock import patch
605+
from databricks.sql.exc import InternalError
606+
607+
real_convert = ThriftResultSet._convert_columnar_table
608+
609+
def convert_dropping_a_row(self_, table):
610+
return real_convert(self_, table)[:-1]
611+
612+
rs = self.make_columnar_result_set(10, arraysize=4)
613+
with patch.object(
614+
ThriftResultSet,
615+
"_convert_columnar_table",
616+
autospec=True,
617+
side_effect=convert_dropping_a_row,
618+
):
619+
with self.assertRaises(InternalError):
620+
rs.fetchone()
621+
582622
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")
583623
def test_fetchone_then_fetchall_arrow_keeps_order(self):
584624
rs = self.make_columnar_result_set(10, arraysize=4)

0 commit comments

Comments
 (0)