Skip to content

Commit 1a44df9

Browse files
ai: apply changes for #1007 (2 review threads)
Addresses: - #4418631883 at src/databricks/sql/backend/sea/result_set.py:250 - #4418631888 at src/databricks/sql/result_set.py:222 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent 5f45f56 commit 1a44df9

4 files changed

Lines changed: 146 additions & 28 deletions

File tree

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

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory
2020
from databricks.sql.backend.types import ExecuteResponse
2121
from databricks.sql.result_set import ResultSet
22-
from databricks.sql.utils import concat_table_chunks
22+
from databricks.sql.utils import CloudFetchQueue, concat_table_chunks
2323

2424
logger = logging.getLogger(__name__)
2525

@@ -186,6 +186,22 @@ def _fetch_raw_from_queue(self, size: Optional[int]) -> Any:
186186
self._next_row_index += self._raw_num_rows(results)
187187
return results
188188

189+
def _fetch_available_from_queue(self, size: int) -> Any:
190+
"""Fetch up to ``size`` rows the queue already holds, for Arrow fetchone() refills.
191+
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's ``_fetch_available_from_queue`` and the kernel's ``_fetch_available``).
197+
Only called while the fetchone() row buffer is empty.
198+
"""
199+
if isinstance(self.results, CloudFetchQueue) and self.results.table:
200+
available = self.results.table.num_rows - self.results.table_row_index
201+
if available > 0:
202+
size = min(size, available)
203+
return self._fetch_raw_from_queue(size)
204+
189205
def _fetch_raw_with_buffer(self, size: Optional[int]) -> Any:
190206
"""Fetch ``size`` (all if None) raw rows, draining the fetchone() buffer first."""
191207
buffered = self._row_buffer_take_raw(size)
@@ -247,7 +263,9 @@ def fetchone(self) -> Optional[Row]:
247263

248264
if isinstance(self.results, JsonQueue):
249265
return self._fetchone_buffered(self.fetchmany_json, self._create_json_table)
250-
return self._fetchone_buffered(self.fetchmany_arrow, self._convert_arrow_table)
266+
return self._fetchone_buffered(
267+
self._fetch_available_from_queue, self._convert_arrow_table
268+
)
251269

252270
def fetchmany(self, size: int) -> List[Row]:
253271
"""

‎src/databricks/sql/result_set.py‎

Lines changed: 56 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,15 @@
3434
_FETCHONE_BUFFER_MAX_ROWS = 1000
3535

3636

37+
class _FailedRow:
38+
"""Placeholder in the fetchone() row buffer for a row that failed to convert."""
39+
40+
__slots__ = ("error",)
41+
42+
def __init__(self, error: Exception):
43+
self.error = error
44+
45+
3746
class ResultSet(ABC):
3847
"""
3948
Abstract base class for result sets returned by different backend implementations.
@@ -179,7 +188,44 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
179188
start, length = self._row_buffer_advance(n)
180189
if self._row_buffer_count() == 0:
181190
self._row_buffer_clear()
182-
return rows[start : start + length]
191+
taken = rows[start : start + length]
192+
for row in taken:
193+
if isinstance(row, _FailedRow):
194+
# The taken rows are consumed either way, matching how a failed
195+
# conversion behaved before fetchone() buffering.
196+
raise row.error
197+
return taken
198+
199+
@staticmethod
200+
def _convert_rows_individually(
201+
raw: Any, convert: Callable[[Any], List[Row]]
202+
) -> Optional[List[Any]]:
203+
"""Convert ``raw`` row by row after a batch conversion failed.
204+
205+
Rows that fail are returned as ``_FailedRow`` placeholders so the error is
206+
raised only when that row is handed out. Returns None if ``raw`` has no rows
207+
to split, in which case the caller re-raises the batch error.
208+
"""
209+
if raw is None:
210+
return None
211+
n = len(raw) if isinstance(raw, list) else getattr(raw, "num_rows", 0)
212+
if not n:
213+
return None
214+
rows: List[Any] = []
215+
for i in range(n):
216+
one = raw[i : i + 1] if isinstance(raw, list) else raw.slice(i, 1)
217+
try:
218+
rows.extend(convert(one))
219+
except Exception as e:
220+
rows.append(_FailedRow(e))
221+
failed = sum(isinstance(r, _FailedRow) for r in rows)
222+
logger.warning(
223+
"fetchone(): failed to convert %d of %d row(s) in a batch; each failing "
224+
"row raises when fetched",
225+
failed,
226+
n,
227+
)
228+
return rows
183229

184230
def _row_buffer_take_raw(self, n: Optional[int] = None) -> Any:
185231
"""Pop up to ``n`` (all if None) un-converted rows from the buffer.
@@ -215,29 +261,18 @@ def _fetchone_buffered(
215261
"""
216262
if self._row_buffer_count() == 0:
217263
raw = fetch_raw(max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1))
218-
# If convert() raises, this batch is dropped on purpose: the queue has
219-
# already moved past it, and _next_row_index (advanced by fetch_raw)
220-
# still equals the queue position. Rolling the index back would make it
221-
# lag the queue and break Thrift's expected_row_start_offset check on
222-
# the next fetch, and retrying convert() on the same data would just
223-
# fail again. A warning records how many rows were discarded so a
224-
# caller that swallows the error and keeps fetching isn't silently
225-
# short of rows.
226264
try:
227265
rows = convert(raw)
228266
except Exception:
229-
if raw is None:
230-
dropped = 0
231-
elif isinstance(raw, list):
232-
dropped = len(raw)
233-
else:
234-
dropped = getattr(raw, "num_rows", 0)
235-
logger.warning(
236-
"fetchone(): failed to convert a batch of %d row(s); these rows "
237-
"are discarded and will not be returned by later fetches",
238-
dropped,
239-
)
240-
raise
267+
# The queue has already moved past this batch, so it can't be
268+
# re-read. Fall back to converting it one row at a time so that,
269+
# as before batching, only the row(s) that fail to convert are
270+
# lost: each failing row is parked as a _FailedRow and re-raises
271+
# its own error when it is reached.
272+
fallback = self._convert_rows_individually(raw, convert)
273+
if fallback is None:
274+
raise
275+
rows = fallback
241276
if not rows:
242277
return None
243278
self._row_buffer_raw = raw

‎tests/unit/test_fetches.py‎

Lines changed: 33 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -534,21 +534,49 @@ def test_mixed_fetches_keep_order_and_rownumber(self):
534534
self.assertEqual(rs.rownumber, 10)
535535
self.assertIsNone(rs.fetchone())
536536

537-
def test_fetchone_convert_failure_logs_dropped_rows_and_reraises(self):
537+
def test_fetchone_convert_failure_loses_only_the_bad_row(self):
538538
from unittest.mock import patch
539539

540+
real_convert = ThriftResultSet._convert_columnar_table
541+
542+
def convert_failing_on_2(self_, table):
543+
rows = real_convert(self_, table)
544+
if any(r[0] == 2 for r in rows):
545+
raise ValueError("bad value")
546+
return rows
547+
540548
rs = self.make_columnar_result_set(10, arraysize=4)
541549
with patch.object(
542550
ThriftResultSet,
543551
"_convert_columnar_table",
544-
side_effect=ValueError("bad value"),
552+
autospec=True,
553+
side_effect=convert_failing_on_2,
545554
), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
555+
self.assertEqual(rs.fetchone()[0], 0)
556+
self.assertEqual(rs.fetchone()[0], 1)
546557
with self.assertRaises(ValueError):
547558
rs.fetchone()
548-
self.assertTrue(any("batch of 4 row(s)" in m for m in logs.output))
549-
# The failed batch is dropped; the next fetch resumes after it.
559+
# Only the failing row is lost; the rest of the batch is still served.
560+
self.assertEqual(rs.rownumber, 3)
561+
self.assertEqual(rs.fetchone()[0], 3)
562+
self.assertEqual([r[0] for r in rs.fetchall()], list(range(4, 10)))
563+
self.assertTrue(any("1 of 4 row(s)" in m for m in logs.output))
564+
self.assertEqual(rs.rownumber, 10)
565+
566+
def test_fetchone_convert_failure_on_every_row_reraises_per_row(self):
567+
from unittest.mock import patch
568+
569+
rs = self.make_columnar_result_set(6, arraysize=4)
570+
with patch.object(
571+
ThriftResultSet,
572+
"_convert_columnar_table",
573+
side_effect=ValueError("bad value"),
574+
):
575+
for expected_rownumber in range(1, 5):
576+
with self.assertRaises(ValueError):
577+
rs.fetchone()
578+
self.assertEqual(rs.rownumber, expected_rownumber)
550579
self.assertEqual(rs._row_buffer_count(), 0)
551-
self.assertEqual(rs.rownumber, 4)
552580
self.assertEqual(rs.fetchone()[0], 4)
553581

554582
@pytest.mark.skipif(pa is None, reason="PyArrow is not installed")

‎tests/unit/test_sea_result_set.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -658,6 +658,43 @@ def test_arrow_fetchone_then_fetchall_arrow_keeps_order(
658658
assert table.column(1).to_pylist() == list(range(1, 10))
659659
assert rs.rownumber == 10
660660

661+
@pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
662+
def test_cloudfetch_fetchone_refill_capped_to_current_chunk(
663+
self, mock_connection, mock_sea_client, execute_response
664+
):
665+
"""fetchone() must not block on the next link while the current chunk has rows."""
666+
from databricks.sql.backend.sea.queue import SeaCloudFetchQueue
667+
668+
chunk = pyarrow.Table.from_pydict(
669+
{
670+
"col1": ["v{}".format(i) for i in range(10)],
671+
"col2": list(range(10)),
672+
"col3": [True] * 10,
673+
}
674+
)
675+
queue = Mock(spec=SeaCloudFetchQueue)
676+
queue.table = chunk
677+
queue.table_row_index = 8
678+
queue.next_n_rows.side_effect = lambda n: chunk.slice(8, n)
679+
result_data = ResultData(data=None, external_links=[], row_count=10)
680+
with patch(
681+
"databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue",
682+
return_value=queue,
683+
):
684+
rs = SeaResultSet(
685+
connection=mock_connection,
686+
execute_response=execute_response,
687+
sea_client=mock_sea_client,
688+
result_data=result_data,
689+
manifest=self._create_empty_manifest(ResultFormat.ARROW_STREAM),
690+
buffer_size_bytes=1000,
691+
arraysize=100,
692+
)
693+
assert rs.fetchone().col2 == 8
694+
queue.next_n_rows.assert_called_once_with(2)
695+
assert rs.fetchone().col2 == 9
696+
assert rs.rownumber == 2
697+
661698
def test_is_staging_operation(
662699
self, mock_connection, mock_sea_client, execute_response
663700
):

0 commit comments

Comments
 (0)