Skip to content

Commit c848f81

Browse files
ai: apply changes for #1007 (3 review threads)
Addresses: - #4419116506 at src/databricks/sql/result_set.py:35 - #4419116511 at src/databricks/sql/result_set.py:448 - #4419116514 at src/databricks/sql/result_set.py:358 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent 86b92d3 commit c848f81

3 files changed

Lines changed: 132 additions & 8 deletions

File tree

‎src/databricks/sql/result_set.py‎

Lines changed: 35 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,11 +29,18 @@
2929
logger = logging.getLogger(__name__)
3030

3131
# Upper bound on how many rows a fetchone() refill pulls and converts at once. The
32-
# refill uses ``arraysize`` but is capped here so the first fetchone() on a cursor
33-
# with a large arraysize (default 100000) doesn't materialise a full page just to
34-
# return one row.
32+
# refill uses ``arraysize`` but is capped here so a cursor with a large arraysize
33+
# (default 100000) doesn't materialise a full page just to return one row.
3534
_FETCHONE_BUFFER_MAX_ROWS = 1000
3635

36+
# fetchone() refills grow geometrically from this size (1 -> 10 -> 100 -> ...) up to
37+
# min(arraysize, _FETCHONE_BUFFER_MAX_ROWS). The first refill converts a single row,
38+
# so a lone fetchone() (e.g. SQLAlchemy ``.first()`` / ``.scalar()`` or an EXISTS
39+
# probe) costs the same as before batching, while iteration loops quickly reach the
40+
# full batch size.
41+
_FETCHONE_BUFFER_INITIAL_ROWS = 1
42+
_FETCHONE_BUFFER_GROWTH_FACTOR = 10
43+
3744
# After a fetchone() batch conversion fails, the batch is retried one row at a time.
3845
# If this many leading rows all fail too, with the same error type and message as
3946
# the batch, the failure is treated as systematic: the rest of the batch is not
@@ -124,6 +131,9 @@ def __init__(
124131
self._row_buffer_pos = 0
125132
# Set once a fetchone() conversion failure has been logged at WARNING.
126133
self._fetchone_conversion_warned = False
134+
# Size of the next fetchone() refill; grows on each refill (see
135+
# ``_FETCHONE_BUFFER_GROWTH_FACTOR``).
136+
self._fetchone_refill_rows = _FETCHONE_BUFFER_INITIAL_ROWS
127137

128138
def __iter__(self):
129139
while True:
@@ -336,8 +346,11 @@ def _fetchone_buffered(
336346
) -> Optional[Row]:
337347
"""Serve fetchone() from the row buffer, refilling in small batches.
338348
339-
Each refill asks ``fetch_raw`` for up to ``arraysize`` rows, capped at
340-
``_FETCHONE_BUFFER_MAX_ROWS``. ``fetch_raw`` may return fewer (Thrift and the
349+
Refill sizes start at ``_FETCHONE_BUFFER_INITIAL_ROWS`` and grow by
350+
``_FETCHONE_BUFFER_GROWTH_FACTOR`` on each refill, capped at ``arraysize``
351+
and ``_FETCHONE_BUFFER_MAX_ROWS``, so a single fetchone() converts only one
352+
row while iteration settles into full batches. ``fetch_raw`` may return
353+
fewer (Thrift and the
341354
kernel backend only serve what is already buffered, refilling at most until
342355
one row is available, so latency to the first row is unchanged).
343356
@@ -346,9 +359,22 @@ def _fetchone_buffered(
346359
queue is only read again once the buffer is empty, so at that point the
347360
index equals the queue position again (Thrift relies on that for
348361
``expected_row_start_offset``).
362+
363+
If conversion yields a different number of rows than the raw chunk held,
364+
``InternalError`` is raised and that whole chunk is dropped (the queue and
365+
``rownumber`` have already moved past it). The result set should be
366+
considered unusable after that error.
349367
"""
350368
if self._row_buffer_count() == 0:
351-
raw = fetch_raw(max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1))
369+
cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1)
370+
target = getattr(
371+
self, "_fetchone_refill_rows", _FETCHONE_BUFFER_INITIAL_ROWS
372+
)
373+
size = max(min(target, cap), 1)
374+
self._fetchone_refill_rows = min(
375+
size * _FETCHONE_BUFFER_GROWTH_FACTOR, _FETCHONE_BUFFER_MAX_ROWS
376+
)
377+
raw = fetch_raw(size)
352378
try:
353379
rows = convert(raw)
354380
except Exception as batch_error:
@@ -379,6 +405,9 @@ def _fetchone_buffered(
379405
# fetch_raw advanced _next_row_index (and the Thrift fetch offset)
380406
# by raw_count; the buffer must line up 1:1 with the raw chunk or
381407
# rownumber and fetchmany_arrow slicing would silently drift.
408+
# The raw chunk is not buffered, so its raw_count rows are lost
409+
# (yet counted in rownumber): treat the result set as unusable
410+
# after this error; later fetches resume from the next chunk.
382411
raise InternalError(
383412
"fetchone(): converted {} row(s) from a chunk of {} raw row(s)".format(
384413
len(rows), raw_count

‎tests/unit/test_fetches.py‎

Lines changed: 70 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -370,6 +370,25 @@ def test_mixed_fetches_across_pages_keep_order_and_refill_offsets(self):
370370
self.assertEqual(rs.rownumber, 10)
371371
self.assertEqual(offsets, [0, 3, 6, 9])
372372

373+
def test_close_drops_fetchone_row_buffer(self):
374+
# Rows parked in the fetchone() buffer must not survive close(). The warm
375+
# refill pulls every row into the buffer, so the queue itself is empty and
376+
# any row returned after close() could only come from a stale buffer.
377+
from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
378+
379+
for fetch_after_close in (
380+
lambda rs: rs.fetchall(),
381+
lambda rs: rs.fetchmany_arrow(10).to_pylist(),
382+
):
383+
rs = self.make_dummy_result_set_from_initial_results([[1], [2], [3], [4]])
384+
rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
385+
self.assertSequenceEqual(rs.fetchone(), [1])
386+
self.assertEqual(rs._row_buffer_count(), 3)
387+
rs.close()
388+
self.assertEqual(rs._row_buffer_count(), 0)
389+
self.assertEqual(fetch_after_close(rs), [])
390+
self.assertIsNone(rs.fetchone())
391+
373392
def test_fetchone_refill_does_not_fetch_extra_pages(self):
374393
# A refill must not make extra round-trips to fill ``arraysize`` when
375394
# server pages are smaller (e.g. byte-limited): it serves what the current
@@ -408,6 +427,10 @@ def test_cloudfetch_fetchone_refill_capped_to_current_chunk(self):
408427
description=[("col0", "integer", None, None, None, None, None)],
409428
)
410429
rs.arraysize = 100
430+
# Skip the adaptive refill ramp so the chunk cap, not the ramp, bounds it.
431+
from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
432+
433+
rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
411434

412435
self.assertSequenceEqual(rs.fetchone(), [8])
413436
queue.next_n_rows.assert_called_once_with(2)
@@ -482,13 +505,21 @@ class ColumnarFetchoneBufferTests(unittest.TestCase):
482505
rows in batches (not one row at a time) while keeping order and rownumber."""
483506

484507
@staticmethod
485-
def make_columnar_result_set(n_rows, arraysize):
508+
def make_columnar_result_set(n_rows, arraysize, warm=True):
509+
"""Build a result set over ``n_rows`` columnar rows.
510+
511+
With ``warm`` (the default) the adaptive fetchone() refill ramp is skipped,
512+
so every refill is ``min(arraysize, _FETCHONE_BUFFER_MAX_ROWS)`` rows and the
513+
batch-boundary assertions below stay simple. Pass ``warm=False`` to test
514+
the ramp itself.
515+
"""
516+
from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
486517
from databricks.sql.utils import ColumnQueue, ColumnTable
487518

488519
table = ColumnTable([list(range(n_rows))], ["col0"])
489520
mock_thrift_backend = Mock(spec=ThriftDatabricksClient)
490521
mock_thrift_backend.fetch_results.return_value = (ColumnQueue(table), False, 0)
491-
return ThriftResultSet(
522+
rs = ThriftResultSet(
492523
connection=Mock(),
493524
execute_response=ExecuteResponse(
494525
command_id=None,
@@ -501,6 +532,9 @@ def make_columnar_result_set(n_rows, arraysize):
501532
thrift_client=mock_thrift_backend,
502533
arraysize=arraysize,
503534
)
535+
if warm:
536+
rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
537+
return rs
504538

505539
@staticmethod
506540
def _count_conversions():
@@ -543,6 +577,40 @@ def test_fetchone_refill_is_capped_for_large_arraysize(self):
543577
self.assertEqual([r[0] for r in rs.fetchall()], list(range(1, n_rows)))
544578
self.assertEqual(rs.rownumber, n_rows)
545579

580+
def test_first_fetchone_converts_a_single_row(self):
581+
# A lone fetchone() (SQLAlchemy .first()/.scalar(), EXISTS probes) must not
582+
# pull and convert a whole batch: the first refill is a single row.
583+
rs = self.make_columnar_result_set(5000, arraysize=100000, warm=False)
584+
with self._count_conversions() as m:
585+
self.assertEqual(rs.fetchone()[0], 0)
586+
self.assertEqual(m.call_count, 1)
587+
self.assertEqual(m.call_args[0][1].num_rows, 1)
588+
self.assertEqual(rs._row_buffer_count(), 0)
589+
self.assertEqual(rs.rownumber, 1)
590+
591+
def test_fetchone_refill_size_ramps_up_to_cap(self):
592+
from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
593+
594+
n_rows = 1 + 10 + 100 + _FETCHONE_BUFFER_MAX_ROWS * 2
595+
rs = self.make_columnar_result_set(n_rows, arraysize=100000, warm=False)
596+
with self._count_conversions() as m:
597+
rows = [r[0] for r in rs]
598+
self.assertEqual(rows, list(range(n_rows)))
599+
sizes = [c[0][1].num_rows for c in m.call_args_list]
600+
self.assertEqual(
601+
sizes[:5],
602+
[1, 10, 100, _FETCHONE_BUFFER_MAX_ROWS, _FETCHONE_BUFFER_MAX_ROWS],
603+
)
604+
self.assertEqual(rs.rownumber, n_rows)
605+
606+
def test_fetchone_refill_ramp_is_capped_by_arraysize(self):
607+
rs = self.make_columnar_result_set(10, arraysize=4, warm=False)
608+
with self._count_conversions() as m:
609+
rows = [r[0] for r in rs]
610+
self.assertEqual(rows, list(range(10)))
611+
sizes = [c[0][1].num_rows for c in m.call_args_list]
612+
self.assertEqual(sizes[:4], [1, 4, 4, 1])
613+
546614
def test_mixed_fetches_keep_order_and_rownumber(self):
547615
rs = self.make_columnar_result_set(10, arraysize=4)
548616
self.assertEqual(rs.fetchone()[0], 0)

‎tests/unit/test_sea_result_set.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,29 @@ def test_close_when_connection_closed(
308308
assert result_set.has_been_closed_server_side is True
309309
assert result_set.status == CommandState.CLOSED
310310

311+
@pytest.mark.parametrize("fetch_after_close", ["fetchall", "fetchmany_arrow"])
312+
def test_close_drops_fetchone_row_buffer(
313+
self, result_set_with_data, fetch_after_close
314+
):
315+
"""Rows parked in the fetchone() buffer must not survive close()."""
316+
if fetch_after_close == "fetchmany_arrow" and pyarrow is None:
317+
pytest.skip("PyArrow is not installed")
318+
from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
319+
320+
rs = result_set_with_data
321+
# Warm refill pulls all 5 rows into the buffer, emptying the queue, so any
322+
# row returned after close() could only come from a stale buffer.
323+
rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
324+
assert rs.fetchone()[0] == "value1"
325+
assert rs._row_buffer_count() == 4
326+
rs.close()
327+
assert rs._row_buffer_count() == 0
328+
if fetch_after_close == "fetchall":
329+
assert rs.fetchall() == []
330+
else:
331+
assert rs.fetchmany_arrow(10).num_rows == 0
332+
assert rs.fetchone() is None
333+
311334
def test_convert_json_types(self, result_set_with_data, sample_data):
312335
"""Test the _convert_json_types method."""
313336
# Call _convert_json_types
@@ -690,6 +713,10 @@ def test_cloudfetch_fetchone_refill_capped_to_current_chunk(
690713
buffer_size_bytes=1000,
691714
arraysize=100,
692715
)
716+
# Skip the adaptive refill ramp so the chunk cap, not the ramp, bounds it.
717+
from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS
718+
719+
rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS
693720
assert rs.fetchone().col2 == 8
694721
queue.next_n_rows.assert_called_once_with(2)
695722
assert rs.fetchone().col2 == 9

0 commit comments

Comments
 (0)