Skip to content

Commit 48a432b

Browse files
ai: apply changes for #1007 (1 review thread)
Addresses: - #4418777541 at src/databricks/sql/result_set.py:276 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent 05f9f56 commit 48a432b

2 files changed

Lines changed: 58 additions & 6 deletions

File tree

‎src/databricks/sql/result_set.py‎

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,12 @@
3434
# return one row.
3535
_FETCHONE_BUFFER_MAX_ROWS = 1000
3636

37+
# After a fetchone() batch conversion fails, the batch is retried one row at a time.
38+
# If this many leading rows all fail too, the failure is treated as systematic: the
39+
# rest of the batch is not retried row by row (each would cost a separate
40+
# conversion) and those rows re-raise the original batch error instead.
41+
_FETCHONE_FALLBACK_PROBE_ROWS = 5
42+
3743

3844
class _FailedRow:
3945
"""Placeholder in the fetchone() row buffer for a row that failed to convert."""
@@ -207,30 +213,44 @@ def _raw_row_count(raw: Any) -> Optional[int]:
207213

208214
@staticmethod
209215
def _convert_rows_individually(
210-
raw: Any, convert: Callable[[Any], List[Row]]
216+
raw: Any, convert: Callable[[Any], List[Row]], batch_error: Exception
211217
) -> Optional[List[Any]]:
212218
"""Convert ``raw`` row by row after a batch conversion failed.
213219
214220
Rows that fail are returned as ``_FailedRow`` placeholders so the error is
215-
raised only when that row is handed out. Returns None if ``raw`` has no rows
216-
to split, in which case the caller re-raises the batch error.
221+
raised only when that row is handed out. If the first
222+
``_FETCHONE_FALLBACK_PROBE_ROWS`` rows all fail, the failure is treated as
223+
systematic and the remaining rows are parked with ``batch_error`` instead of
224+
being converted one by one. Returns None if ``raw`` has no rows to split, in
225+
which case the caller re-raises the batch error.
217226
"""
218227
n = ResultSet._raw_row_count(raw)
219228
if not n:
220229
return None
221230
rows: List[Any] = []
231+
systematic = False
222232
for i in range(n):
233+
if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(
234+
isinstance(r, _FailedRow) for r in rows
235+
):
236+
systematic = True
237+
rows.extend(_FailedRow(batch_error) for _ in range(n - i))
238+
break
223239
one = raw[i : i + 1] if isinstance(raw, list) else raw.slice(i, 1)
224240
try:
225241
rows.extend(convert(one))
226242
except Exception as e:
227243
rows.append(_FailedRow(e))
228244
failed = sum(isinstance(r, _FailedRow) for r in rows)
229245
logger.warning(
230-
"fetchone(): failed to convert %d of %d row(s) in a batch; each failing "
246+
"fetchone(): failed to convert %d of %d row(s) in a batch%s; each failing "
231247
"row raises when fetched",
232248
failed,
233249
n,
250+
" (first %d rows all failed, rest not retried)" % _FETCHONE_FALLBACK_PROBE_ROWS
251+
if systematic
252+
else "",
253+
exc_info=batch_error,
234254
)
235255
return rows
236256

@@ -270,13 +290,13 @@ def _fetchone_buffered(
270290
raw = fetch_raw(max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1))
271291
try:
272292
rows = convert(raw)
273-
except Exception:
293+
except Exception as batch_error:
274294
# The queue has already moved past this batch, so it can't be
275295
# re-read. Fall back to converting it one row at a time so that,
276296
# as before batching, only the row(s) that fail to convert are
277297
# lost: each failing row is parked as a _FailedRow and re-raises
278298
# its own error when it is reached.
279-
fallback = self._convert_rows_individually(raw, convert)
299+
fallback = self._convert_rows_individually(raw, convert, batch_error)
280300
if fallback is None:
281301
raise
282302
rows = fallback

‎tests/unit/test_fetches.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -600,6 +600,38 @@ def test_fetchone_convert_failure_on_every_row_reraises_per_row(self):
600600
self.assertEqual(rs._row_buffer_count(), 0)
601601
self.assertEqual(rs.fetchone()[0], 4)
602602

603+
def test_fetchone_systematic_convert_failure_stops_per_row_retries(self):
604+
from unittest.mock import patch
605+
from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS
606+
607+
batch_error = ValueError("batch failure")
608+
calls = []
609+
610+
def convert(self_, table):
611+
calls.append(table.num_rows)
612+
if len(calls) == 1:
613+
raise batch_error
614+
raise ValueError("row failure")
615+
616+
rs = self.make_columnar_result_set(50, arraysize=20)
617+
with patch.object(
618+
ThriftResultSet,
619+
"_convert_columnar_table",
620+
autospec=True,
621+
side_effect=convert,
622+
), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
623+
for i in range(20):
624+
with self.assertRaises(ValueError) as ctx:
625+
rs.fetchone()
626+
if i >= _FETCHONE_FALLBACK_PROBE_ROWS:
627+
self.assertIs(ctx.exception, batch_error)
628+
self.assertEqual(rs.rownumber, i + 1)
629+
# One batch attempt plus only the probe rows were converted individually.
630+
self.assertEqual(len(calls), 1 + _FETCHONE_FALLBACK_PROBE_ROWS)
631+
self.assertTrue(any("20 of 20 row(s)" in m for m in logs.output))
632+
self.assertTrue(any("batch failure" in m for m in logs.output))
633+
self.assertEqual(rs.fetchone()[0], 20)
634+
603635
def test_fetchone_raw_converted_count_mismatch_fails_loudly(self):
604636
from unittest.mock import patch
605637
from databricks.sql.exc import InternalError

0 commit comments

Comments
 (0)