Skip to content

Commit 2199868

Browse files
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>
1 parent 57fb8a9 commit 2199868

3 files changed

Lines changed: 158 additions & 7 deletions

File tree

‎src/databricks/sql/result_set.py‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -202,8 +202,10 @@ def _fetchone_buffered(
202202
) -> Optional[Row]:
203203
"""Serve fetchone() from the row buffer, refilling in small batches.
204204
205-
Each refill fetches ``arraysize`` rows, capped at
206-
``_FETCHONE_BUFFER_MAX_ROWS``.
205+
Each refill asks ``fetch_raw`` for up to ``arraysize`` rows, capped at
206+
``_FETCHONE_BUFFER_MAX_ROWS``. ``fetch_raw`` may return fewer (Thrift only
207+
serves what its current queue holds, refilling it at most until one row is
208+
available, so latency to the first row is unchanged).
207209
208210
``_next_row_index`` keeps counting rows handed to the caller: rows parked in
209211
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):
484486

485487
return concat_table_chunks(partial_result_chunks)
486488

489+
def _fetch_available_from_queue(self, size: int):
490+
"""Fetch up to ``size`` rows the queue already holds, for fetchone() refills.
491+
492+
Unlike ``_fetchmany_*_from_queue``, this does not keep calling the server
493+
until ``size`` rows are collected: it only refills the queue while it is
494+
empty, so fetchone() makes no more round-trips than it needs to return one
495+
row (server pages can be smaller than ``size`` when they hit the byte limit).
496+
Only called while the fetchone() row buffer is empty.
497+
"""
498+
results = self.results.next_n_rows(size)
499+
while (
500+
results.num_rows == 0
501+
and not self.has_been_closed_server_side
502+
and self.has_more_rows
503+
):
504+
self._fill_results_buffer()
505+
results = self.results.next_n_rows(size)
506+
self._next_row_index += results.num_rows
507+
return results
508+
487509
def fetchall_arrow(self) -> "pyarrow.Table":
488510
"""Fetch all (remaining) rows of a query result, returning them as a PyArrow table."""
489511
buffered = self._row_buffer_take_raw()
@@ -553,9 +575,11 @@ def fetchone(self) -> Optional[Row]:
553575
"""
554576
if isinstance(self.results, ColumnQueue):
555577
return self._fetchone_buffered(
556-
self.fetchmany_columnar, self._convert_columnar_table
578+
self._fetch_available_from_queue, self._convert_columnar_table
557579
)
558-
return self._fetchone_buffered(self.fetchmany_arrow, self._convert_arrow_table)
580+
return self._fetchone_buffered(
581+
self._fetch_available_from_queue, self._convert_arrow_table
582+
)
559583

560584
def fetchall(self) -> List[Row]:
561585
"""

‎tests/unit/test_fetches.py‎

Lines changed: 65 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,8 @@ def make_dummy_result_set_from_initial_results(initial_results):
9595
return rs
9696

9797
@staticmethod
98-
def make_dummy_result_set_from_batch_list(batch_list):
98+
def make_dummy_result_set_from_batch_list(batch_list, offsets=None):
99+
"""``offsets``, if given, records ``expected_row_start_offset`` per refill."""
99100
batch_index = 0
100101

101102
def fetch_results(
@@ -110,6 +111,8 @@ def fetch_results(
110111
chunk_id=0,
111112
):
112113
nonlocal batch_index
114+
if offsets is not None:
115+
offsets.append(expected_row_start_offset)
113116
results = FetchTests.make_arrow_queue(batch_list[batch_index])
114117
batch_index += 1
115118

@@ -328,6 +331,67 @@ def test_fetchone_without_initial_results(self):
328331
dummy_result_set = self.make_dummy_result_set_from_batch_list(batch_list_2)
329332
self.assertEqual(dummy_result_set.fetchone(), None)
330333

334+
def test_mixed_fetches_across_pages_keep_order_and_refill_offsets(self):
335+
# Issue #981: rows parked in the fetchone() buffer must not skew
336+
# ``_next_row_index``, which Thrift sends as ``expected_row_start_offset``
337+
# on every refill. Pages are 3 rows; fetchone() refills 2 rows at a time,
338+
# so buffer refills and server page boundaries interleave.
339+
batch_list = [[[0], [1], [2]], [[3], [4], [5]], [[6], [7], [8]], [[9]]]
340+
offsets = []
341+
rs = self.make_dummy_result_set_from_batch_list(batch_list, offsets=offsets)
342+
rs.arraysize = 2
343+
self.assertEqual(offsets, [0])
344+
345+
self.assertSequenceEqual(rs.fetchone(), [0])
346+
self.assertEqual(rs.rownumber, 1)
347+
348+
# Drains the buffered row, the rest of page 0, then refills page 1.
349+
table = rs.fetchmany_arrow(3)
350+
self.assertEqual(table.column(0).to_pylist(), [1, 2, 3])
351+
self.assertEqual(rs.rownumber, 4)
352+
self.assertEqual(offsets, [0, 3])
353+
354+
self.assertSequenceEqual(rs.fetchone(), [4])
355+
self.assertSequenceEqual(rs.fetchone(), [5])
356+
self.assertEqual(rs.rownumber, 6)
357+
self.assertEqual(offsets, [0, 3])
358+
359+
# Buffer empty and page 1 exhausted: the refill crosses into page 2.
360+
self.assertSequenceEqual(rs.fetchone(), [6])
361+
self.assertEqual(rs.rownumber, 7)
362+
self.assertEqual(offsets, [0, 3, 6])
363+
364+
# Drains the buffered row [7], the rest of page 2, then page 3.
365+
self.assertEqualRowValues(rs.fetchall(), [[7], [8], [9]])
366+
self.assertEqual(rs.rownumber, 10)
367+
self.assertEqual(offsets, [0, 3, 6, 9])
368+
369+
self.assertIsNone(rs.fetchone())
370+
self.assertEqual(rs.rownumber, 10)
371+
self.assertEqual(offsets, [0, 3, 6, 9])
372+
373+
def test_fetchone_refill_does_not_fetch_extra_pages(self):
374+
# A refill must not make extra round-trips to fill ``arraysize`` when
375+
# server pages are smaller (e.g. byte-limited): it serves what the current
376+
# queue holds and only fetches a new page once the queue is empty.
377+
batch_list = [[[0], [1]], [[2], [3]], [[4]]]
378+
offsets = []
379+
rs = self.make_dummy_result_set_from_batch_list(batch_list, offsets=offsets)
380+
rs.arraysize = 100
381+
self.assertEqual(offsets, [0])
382+
383+
self.assertSequenceEqual(rs.fetchone(), [0])
384+
self.assertEqual(offsets, [0])
385+
self.assertSequenceEqual(rs.fetchone(), [1])
386+
self.assertEqual(offsets, [0])
387+
self.assertSequenceEqual(rs.fetchone(), [2])
388+
self.assertEqual(offsets, [0, 2])
389+
self.assertEqual(rs.rownumber, 3)
390+
self.assertEqualRowValues(rs.fetchall(), [[3], [4]])
391+
self.assertEqual(offsets, [0, 2, 4])
392+
self.assertIsNone(rs.fetchone())
393+
self.assertEqual(rs.rownumber, 5)
394+
331395
# Regression tests for fetchmany_arrow / fetchall_arrow handling of
332396
# the schemaless CloudFetch placeholder.
333397
def test_fetchall_arrow_drops_mismatched_empty_placeholder(self):

‎tests/unit/test_kernel_result_set.py‎

Lines changed: 65 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ def close(self):
4343
self.closed = True
4444

4545

46-
def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet:
46+
def _make_rs(handle, *, disable_pandas=False, arraysize=100) -> KernelResultSet:
4747
# The base ResultSet __init__ takes a `connection` ref it never
4848
# actually dereferences during these buffer tests, so a Mock is
4949
# fine.
@@ -55,7 +55,7 @@ def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet:
5555
backend=backend,
5656
kernel_handle=handle,
5757
command_id=CommandId.from_sea_statement_id("smoke-test"),
58-
arraysize=100,
58+
arraysize=arraysize,
5959
buffer_size_bytes=1024,
6060
)
6161

@@ -184,6 +184,69 @@ def test_fetchall_rows(int_schema):
184184
assert [r[0] for r in rows] == [1, 2, 3]
185185

186186

187+
def _ten_rows_handle(int_schema):
188+
# Batch boundaries (3/4/3) deliberately don't line up with arraysize=4
189+
# so fetchone() refills span kernel batches.
190+
return _FakeKernelHandle(
191+
int_schema,
192+
[
193+
_batch(int_schema, [0, 1, 2]),
194+
_batch(int_schema, [3, 4, 5, 6]),
195+
_batch(int_schema, [7, 8, 9]),
196+
],
197+
)
198+
199+
200+
def test_mixed_fetches_keep_order_and_rownumber(int_schema):
201+
rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
202+
assert rs.fetchone()[0] == 0
203+
assert rs.fetchone()[0] == 1
204+
assert rs.rownumber == 2
205+
assert [r[0] for r in rs.fetchmany(3)] == [2, 3, 4]
206+
assert rs.rownumber == 5
207+
assert rs.fetchone()[0] == 5
208+
assert rs.rownumber == 6
209+
assert [r[0] for r in rs.fetchall()] == [6, 7, 8, 9]
210+
assert rs.rownumber == 10
211+
assert rs.fetchone() is None
212+
assert rs.rownumber == 10
213+
214+
215+
def test_fetchone_then_fetchall_arrow_keeps_order(int_schema):
216+
rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
217+
assert rs.fetchone()[0] == 0
218+
assert rs.rownumber == 1
219+
table = rs.fetchall_arrow()
220+
assert table.column(0).to_pylist() == list(range(1, 10))
221+
assert rs.rownumber == 10
222+
223+
224+
def test_fetchone_then_fetchmany_arrow_keeps_order(int_schema):
225+
rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
226+
assert rs.fetchone()[0] == 0
227+
# Served entirely from the fetchone() row buffer (3 rows left in it).
228+
assert rs.fetchmany_arrow(2).column(0).to_pylist() == [1, 2]
229+
assert rs.rownumber == 3
230+
# Drains the last buffered row, then continues from the batch FIFO.
231+
assert rs.fetchmany_arrow(3).column(0).to_pylist() == [3, 4, 5]
232+
assert rs.rownumber == 6
233+
assert rs.fetchmany_arrow(0).num_rows == 0
234+
assert rs.rownumber == 6
235+
assert rs.fetchmany_arrow(10).column(0).to_pylist() == [6, 7, 8, 9]
236+
assert rs.rownumber == 10
237+
238+
239+
def test_fetchone_then_fetchmany_rows_spans_buffer(int_schema):
240+
rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
241+
assert rs.fetchone()[0] == 0
242+
# 3 rows come from the row buffer, the rest from the batch FIFO.
243+
assert [r[0] for r in rs.fetchmany(6)] == [1, 2, 3, 4, 5, 6]
244+
assert rs.rownumber == 7
245+
assert [r[0] for r in rs.fetchmany(10)] == [7, 8, 9]
246+
assert rs.rownumber == 10
247+
assert rs.fetchmany(1) == []
248+
249+
187250
def test_fetchmany_negative_raises(int_schema):
188251
rs = _make_rs(_FakeKernelHandle(int_schema, []))
189252
with pytest.raises(ValueError):

0 commit comments

Comments
 (0)