Skip to content

Commit 9e4320f

Browse files
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>
1 parent 4b590fa commit 9e4320f

4 files changed

Lines changed: 17 additions & 13 deletions

File tree

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -189,12 +189,11 @@ def _fetch_raw_from_queue(self, size: Optional[int]) -> Any:
189189
def _fetch_available_from_queue(self, size: int) -> Any:
190190
"""Fetch up to ``size`` rows the queue already holds, for Arrow fetchone() refills.
191191
192-
On CloudFetch, ``next_n_rows`` keeps walking downloaded links until it has
193-
``size`` rows, so a refill that starts near the end of the current chunk
194-
would block on the next download before returning a row that is already
195-
available. Cap the request at what the current chunk still holds (matching
196-
Thrift and the kernel's ``_fetch_available``) via the shared
197-
``ResultSet._cap_to_current_chunk``.
192+
On CloudFetch, the request is capped at what the current chunk still holds
193+
(matching Thrift) via the shared ``ResultSet._cap_to_current_chunk``, so a
194+
refill is limited to one chunk rather than stitching several links into one
195+
batch. It does not avoid waiting on the next download: reading to the end of
196+
a chunk still makes ``next_n_rows`` advance to the next link in that call.
198197
Only called while the fetchone() row buffer is empty.
199198
"""
200199
return self._fetch_raw_from_queue(self._cap_to_current_chunk(size))

‎src/databricks/sql/result_set.py‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,9 @@ def _convert_rows_individually(
257257
def _cap_to_current_chunk(self, size: int) -> int:
258258
"""Cap ``size`` at the rows left in the current CloudFetch chunk, if any.
259259
260+
Limits each fetchone() refill to a single CloudFetch chunk. It does not
261+
avoid waiting on the next download: ``next_n_rows`` advances to the next
262+
file as soon as the current one is fully read, inside the same call.
260263
Shared by the Thrift and SEA fetchone() refills so the CloudFetchQueue
261264
internals (``table`` / ``table_row_index``) are read in one place.
262265
"""
@@ -604,10 +607,11 @@ def _fetch_available_from_queue(self, size: int):
604607
until ``size`` rows are collected: it only refills the queue while it is
605608
empty, so fetchone() makes no more round-trips than it needs to return one
606609
row (server pages can be smaller than ``size`` when they hit the byte limit).
607-
On CloudFetch, ``next_n_rows`` keeps walking downloaded files until it has
608-
``size`` rows, so the request is also capped at what the current chunk still
609-
holds; otherwise a refill near a chunk boundary would block on the next
610-
download before returning a row that is already available.
610+
On CloudFetch, the request is also capped at what the current chunk still
611+
holds (see ``_cap_to_current_chunk``), so a refill never stitches rows from
612+
several downloaded files into one batch. This only bounds the batch: reading
613+
to the end of a chunk still makes ``next_n_rows`` advance to (and wait for)
614+
the next file before it returns, exactly as the uncapped path did.
611615
Only called while the fetchone() row buffer is empty.
612616
"""
613617
results = self.results.next_n_rows(self._cap_to_current_chunk(size))

‎tests/unit/test_fetches.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -393,8 +393,9 @@ def test_fetchone_refill_does_not_fetch_extra_pages(self):
393393
self.assertEqual(rs.rownumber, 5)
394394

395395
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.
396+
# The fetchone() refill is capped at the rows left in the current CloudFetch
397+
# chunk, so one refill never spans several files. (This only bounds the
398+
# batch size; the mock queue does not model advancing to the next file.)
398399
from databricks.sql.utils import ThriftCloudFetchQueue
399400

400401
_, chunk = self.make_arrow_table([[i] for i in range(10)])

‎tests/unit/test_sea_result_set.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -662,7 +662,7 @@ def test_arrow_fetchone_then_fetchall_arrow_keeps_order(
662662
def test_cloudfetch_fetchone_refill_capped_to_current_chunk(
663663
self, mock_connection, mock_sea_client, execute_response
664664
):
665-
"""fetchone() must not block on the next link while the current chunk has rows."""
665+
"""fetchone() refill is capped at the rows left in the current chunk (batch size only)."""
666666
from databricks.sql.backend.sea.queue import SeaCloudFetchQueue
667667

668668
chunk = pyarrow.Table.from_pydict(

0 commit comments

Comments
 (0)