Skip to content

Commit af2c1fc

Browse files
ai: apply changes for #1007 (4 review threads)
Addresses: - #4418502676 at src/databricks/sql/backend/kernel/result_set.py:229 - #4418502680 at src/databricks/sql/result_set.py:100 - #4418502685 at src/databricks/sql/result_set.py:414 - #4418502697 at src/databricks/sql/backend/sea/result_set.py:202 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent bde5cf8 commit af2c1fc

4 files changed

Lines changed: 131 additions & 12 deletions

File tree

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

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,17 @@ def _fetchmany_arrow_unbuffered(self, size: int) -> pyarrow.Table:
197197
self._ensure_buffered(size)
198198
return self._take_buffered(size)
199199

200+
def _fetch_available(self, size: int) -> pyarrow.Table:
201+
"""Take up to ``size`` rows already buffered, for fetchone() refills.
202+
203+
Unlike ``_fetchmany_arrow_unbuffered`` this does not keep pulling
204+
batches until ``size`` rows are collected: it only pulls from the
205+
kernel while the batch FIFO is empty, so latency to the first row
206+
matches the pre-buffering behaviour.
207+
"""
208+
self._ensure_buffered(1)
209+
return self._take_buffered(min(size, self._buffered_rows()))
210+
200211
@staticmethod
201212
def _prepend_buffered(
202213
buffered: Optional[pyarrow.Table], rest: pyarrow.Table
@@ -228,7 +239,7 @@ def fetchone(self) -> Optional[Row]:
228239
# Rows are converted ``arraysize`` at a time and served from the
229240
# base-class row buffer; converting one row per call is very slow.
230241
return self._fetchone_buffered(
231-
self._fetchmany_arrow_unbuffered, self._convert_arrow_table
242+
self._fetch_available, self._convert_arrow_table
232243
)
233244

234245
def fetchmany(self, size: int) -> List[Row]:
@@ -258,6 +269,9 @@ def close(self) -> None:
258269
connection-close would otherwise issue a kernel call into an
259270
already-disposed session.
260271
"""
272+
# Drop rows parked by fetchone() so post-close fetches return
273+
# nothing, matching the cleared batch FIFO below.
274+
self._row_buffer_clear()
261275
if self._kernel_handle is None:
262276
return
263277
if not self.connection.open:

‎src/databricks/sql/result_set.py‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -203,9 +203,9 @@ def _fetchone_buffered(
203203
"""Serve fetchone() from the row buffer, refilling in small batches.
204204
205205
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).
206+
``_FETCHONE_BUFFER_MAX_ROWS``. ``fetch_raw`` may return fewer (Thrift and the
207+
kernel backend only serve what is already buffered, refilling at most until
208+
one row is available, so latency to the first row is unchanged).
209209
210210
``_next_row_index`` keeps counting rows handed to the caller: rows parked in
211211
the buffer are subtracted on refill and added back as they are taken. The
@@ -283,6 +283,8 @@ def close(self) -> None:
283283
if isinstance(e.args[1], CursorAlreadyClosedError):
284284
logger.info("Operation was canceled by a prior request")
285285
finally:
286+
# Drop rows parked by fetchone() so post-close fetches return nothing.
287+
self._row_buffer_clear()
286288
self.has_been_closed_server_side = True
287289
self.status = CommandState.CLOSED
288290

@@ -411,7 +413,7 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table":
411413
An empty sequence is returned when no more rows are available.
412414
"""
413415
if size < 0:
414-
raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
416+
raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
415417
buffered = self._row_buffer_take_raw(size)
416418
remaining = size - (buffered.num_rows if buffered is not None else 0)
417419
rest = self._fetchmany_arrow_from_queue(remaining) if remaining > 0 else None
@@ -425,7 +427,7 @@ def fetchmany_columnar(self, size: int):
425427
An empty sequence is returned when no more rows are available.
426428
"""
427429
if size < 0:
428-
raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
430+
raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
429431
buffered = self._row_buffer_take_raw(size)
430432
remaining = size - (buffered.num_rows if buffered is not None else 0)
431433
rest = self._fetchmany_columnar_from_queue(remaining) if remaining > 0 else None
@@ -436,7 +438,7 @@ def fetchmany_columnar(self, size: int):
436438
def _fetchmany_arrow_from_queue(self, size: int) -> "pyarrow.Table":
437439
"""Fetch up to ``size`` rows from the results queue, bypassing the row buffer."""
438440
if size < 0:
439-
raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
441+
raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
440442

441443
# Hold 0-row chunks aside instead of appending them to ``partial_result_chunks``.
442444
# CloudFetchQueue may return a placeholder empty table whose schema does not
@@ -473,7 +475,7 @@ def _fetchmany_arrow_from_queue(self, size: int) -> "pyarrow.Table":
473475
def _fetchmany_columnar_from_queue(self, size: int):
474476
"""Fetch up to ``size`` rows from the results queue, bypassing the row buffer."""
475477
if size < 0:
476-
raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
478+
raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
477479

478480
results = self.results.next_n_rows(size)
479481
n_remaining_rows = size - results.num_rows
@@ -605,7 +607,7 @@ def fetchmany(self, size: int) -> List[Row]:
605607
An empty sequence is returned when no more rows are available.
606608
"""
607609
if size < 0:
608-
raise ValueError("size argument for fetchmany is %s but must be >= 0", size)
610+
raise ValueError(f"size argument for fetchmany is {size} but must be >= 0")
609611
buffered = self._row_buffer_take_rows(size)
610612
remaining = size - len(buffered)
611613
if buffered and remaining == 0:

‎tests/unit/test_kernel_result_set.py‎

Lines changed: 27 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,20 @@ def test_mixed_fetches_keep_order_and_rownumber(int_schema):
212212
assert rs.rownumber == 10
213213

214214

215+
def test_fetchone_refill_pulls_at_most_one_batch(int_schema):
216+
# First-row latency: a fetchone() refill must not loop fetch_next_batch()
217+
# until arraysize rows are buffered.
218+
handle = _ten_rows_handle(int_schema)
219+
rs = _make_rs(handle, arraysize=100)
220+
assert rs.fetchone()[0] == 0
221+
assert len(handle._batches) == 2
222+
assert rs.fetchone()[0] == 1
223+
assert rs.fetchone()[0] == 2
224+
assert len(handle._batches) == 2
225+
assert rs.fetchone()[0] == 3
226+
assert len(handle._batches) == 1
227+
228+
215229
def test_fetchone_then_fetchall_arrow_keeps_order(int_schema):
216230
rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
217231
assert rs.fetchone()[0] == 0
@@ -224,10 +238,11 @@ def test_fetchone_then_fetchall_arrow_keeps_order(int_schema):
224238
def test_fetchone_then_fetchmany_arrow_keeps_order(int_schema):
225239
rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
226240
assert rs.fetchone()[0] == 0
227-
# Served entirely from the fetchone() row buffer (3 rows left in it).
241+
# Served entirely from the fetchone() row buffer (2 rows left in it: the
242+
# refill only took the first kernel batch).
228243
assert rs.fetchmany_arrow(2).column(0).to_pylist() == [1, 2]
229244
assert rs.rownumber == 3
230-
# Drains the last buffered row, then continues from the batch FIFO.
245+
# Row buffer is empty; continues from the batch FIFO.
231246
assert rs.fetchmany_arrow(3).column(0).to_pylist() == [3, 4, 5]
232247
assert rs.rownumber == 6
233248
assert rs.fetchmany_arrow(0).num_rows == 0
@@ -239,7 +254,7 @@ def test_fetchone_then_fetchmany_arrow_keeps_order(int_schema):
239254
def test_fetchone_then_fetchmany_rows_spans_buffer(int_schema):
240255
rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
241256
assert rs.fetchone()[0] == 0
242-
# 3 rows come from the row buffer, the rest from the batch FIFO.
257+
# 2 rows come from the row buffer, the rest from the batch FIFO.
243258
assert [r[0] for r in rs.fetchmany(6)] == [1, 2, 3, 4, 5, 6]
244259
assert rs.rownumber == 7
245260
assert [r[0] for r in rs.fetchmany(10)] == [7, 8, 9]
@@ -264,6 +279,15 @@ def test_close_is_idempotent_and_calls_handle(int_schema):
264279
rs.close() # second call is a no-op (kernel handle is None)
265280

266281

282+
def test_close_drops_fetchone_row_buffer(int_schema):
283+
# Rows parked in the fetchone() buffer must not survive close().
284+
rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4)
285+
assert rs.fetchone()[0] == 0
286+
rs.close()
287+
assert rs.fetchone() is None
288+
assert rs.fetchall() == []
289+
290+
267291
def test_empty_stream(int_schema):
268292
rs = _make_rs(_FakeKernelHandle(int_schema, []))
269293
assert rs.fetchone() is None

‎tests/unit/test_sea_result_set.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -579,6 +579,85 @@ def test_json_fetchone_then_fetchall_arrow_keeps_order(
579579
assert table.column(1).to_pylist() == list(range(1, 10))
580580
assert rs.rownumber == 10
581581

582+
def _make_arrow_result_set(
583+
self, mock_connection, mock_sea_client, execute_response, n_rows, arraysize
584+
):
585+
from databricks.sql.utils import ArrowQueue
586+
587+
mock_connection.disable_pandas = True
588+
table = pyarrow.Table.from_pydict(
589+
{
590+
"col1": ["v{}".format(i) for i in range(n_rows)],
591+
"col2": list(range(n_rows)),
592+
"col3": [True] * n_rows,
593+
}
594+
)
595+
result_data = ResultData(data=None, external_links=None, row_count=n_rows)
596+
with patch(
597+
"databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue",
598+
return_value=ArrowQueue(table, n_rows),
599+
):
600+
return SeaResultSet(
601+
connection=mock_connection,
602+
execute_response=execute_response,
603+
sea_client=mock_sea_client,
604+
result_data=result_data,
605+
manifest=self._create_empty_manifest(ResultFormat.ARROW_STREAM),
606+
buffer_size_bytes=1000,
607+
arraysize=arraysize,
608+
)
609+
610+
@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
611+
def test_arrow_mixed_fetches_keep_order_and_rownumber(
612+
self, mock_connection, mock_sea_client, execute_response
613+
):
614+
"""Buffered Arrow slice joined with the queue's next chunk via concat."""
615+
rs = self._make_arrow_result_set(
616+
mock_connection, mock_sea_client, execute_response, 10, arraysize=4
617+
)
618+
assert rs.fetchone().col2 == 0
619+
assert rs.rownumber == 1
620+
# 3 rows left in the buffer; ask for 5 so 2 must come from the queue.
621+
table = rs.fetchmany_arrow(5)
622+
assert isinstance(table, pyarrow.Table)
623+
assert table.column(1).to_pylist() == [1, 2, 3, 4, 5]
624+
assert rs.rownumber == 6
625+
assert rs.fetchone().col2 == 6
626+
assert rs.rownumber == 7
627+
# Fully served from the buffer (remaining == 0 path).
628+
assert rs.fetchmany_arrow(1).column(1).to_pylist() == [7]
629+
assert rs.rownumber == 8
630+
assert [r.col2 for r in rs.fetchall()] == [8, 9]
631+
assert rs.rownumber == 10
632+
assert rs.fetchone() is None
633+
634+
@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
635+
def test_arrow_buffer_with_exhausted_queue_returns_buffer(
636+
self, mock_connection, mock_sea_client, execute_response
637+
):
638+
"""Queue returns a 0-row slice after the buffer drains: buffered rows only."""
639+
rs = self._make_arrow_result_set(
640+
mock_connection, mock_sea_client, execute_response, 3, arraysize=10
641+
)
642+
assert rs.fetchone().col2 == 0
643+
table = rs.fetchmany_arrow(5)
644+
assert table.column(1).to_pylist() == [1, 2]
645+
assert rs.rownumber == 3
646+
assert rs.fetchall_arrow().num_rows == 0
647+
assert rs.rownumber == 3
648+
649+
@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
650+
def test_arrow_fetchone_then_fetchall_arrow_keeps_order(
651+
self, mock_connection, mock_sea_client, execute_response
652+
):
653+
rs = self._make_arrow_result_set(
654+
mock_connection, mock_sea_client, execute_response, 10, arraysize=4
655+
)
656+
assert rs.fetchone().col2 == 0
657+
table = rs.fetchall_arrow()
658+
assert table.column(1).to_pylist() == list(range(1, 10))
659+
assert rs.rownumber == 10
660+
582661
def test_is_staging_operation(
583662
self, mock_connection, mock_sea_client, execute_response
584663
):

0 commit comments

Comments
 (0)