Skip to content

Commit 77e1090

Browse files
ai: apply changes for #1007 (1 review thread)
Addresses: - #4419207269 at src/databricks/sql/result_set.py:280 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent c848f81 commit 77e1090

2 files changed

Lines changed: 77 additions & 69 deletions

File tree

‎src/databricks/sql/result_set.py‎

Lines changed: 32 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -41,24 +41,15 @@
4141
_FETCHONE_BUFFER_INITIAL_ROWS = 1
4242
_FETCHONE_BUFFER_GROWTH_FACTOR = 10
4343

44-
# After a fetchone() batch conversion fails, the batch is retried one row at a time.
45-
# If this many leading rows all fail too, with the same error type and message as
46-
# the batch, the failure is treated as systematic: the rest of the batch is not
47-
# retried row by row (each would cost a separate conversion) and those rows re-raise
48-
# the original batch error instead. Leading rows failing for other reasons do not
49-
# trigger the shortcut, so later good rows are still converted.
50-
_FETCHONE_FALLBACK_PROBE_ROWS = 5
51-
52-
5344
class _FailedRow:
5445
"""Placeholder in the fetchone() row buffer for a row that failed to convert."""
5546

5647
__slots__ = ("error", "traceback")
5748

5849
def __init__(self, error: Exception):
5950
self.error = error
60-
# Snapshot the conversion traceback. Systematic failures park the same
61-
# exception instance in many rows; restoring this before each raise keeps
51+
# Snapshot the conversion traceback. A converter may raise the same
52+
# exception instance for many rows; restoring this before each raise keeps
6253
# its __traceback__ from growing (and pinning frames) on every fetchone().
6354
self.traceback = error.__traceback__
6455

@@ -250,60 +241,50 @@ def _convert_rows_individually(
250241
batch_error: Exception,
251242
log_level: int = logging.WARNING,
252243
) -> Optional[List[Any]]:
253-
"""Convert ``raw`` row by row after a batch conversion failed.
254-
255-
Rows that fail are returned as ``_FailedRow`` placeholders so the error is
256-
raised only when that row is handed out. If the first
257-
``_FETCHONE_FALLBACK_PROBE_ROWS`` rows all fail with the same error as the
258-
batch (same type and message), the failure is treated as systematic and the
259-
remaining rows are parked with ``batch_error`` instead of being converted one
260-
by one. Otherwise every row is tried individually. Returns None if ``raw``
261-
has no rows to split, in which case the caller re-raises the batch error.
244+
"""Convert ``raw`` in smaller pieces after a batch conversion failed.
245+
246+
The failed batch is bisected: each failing range is split in half and the
247+
halves are converted separately, down to single rows. Only rows that fail
248+
on their own are returned as ``_FailedRow`` placeholders (raised when that
249+
row is handed out); every row that converts is kept, so a failure is never
250+
attributed to a row that was not itself tried. Cost is about
251+
``2 * k * log2(n)`` conversions for ``k`` bad rows, and at most ``2n``
252+
when every row fails. Returns None if ``raw`` has no rows to split, in
253+
which case the caller re-raises the batch error.
262254
"""
263255
n = ResultSet._raw_row_count(raw)
264256
if not n:
265257
return None
266258

267-
def same_as_batch(r: Any) -> bool:
268-
return (
269-
isinstance(r, _FailedRow)
270-
and type(r.error) is type(batch_error)
271-
and str(r.error) == str(batch_error)
272-
)
273-
274259
rows: List[Any] = []
275-
systematic = False
276-
for i in range(n):
277-
# Only skip the remaining per-row retries when every probe row failed
278-
# with the same error as the batch; leading rows that fail for their
279-
# own reasons must not cause later good rows to be dropped.
280-
if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(
281-
same_as_batch(r) for r in rows
282-
):
283-
systematic = True
284-
rows.extend(_FailedRow(batch_error) for _ in range(n - i))
285-
break
286-
one = raw[i : i + 1] if isinstance(raw, list) else raw.slice(i, 1)
287-
try:
288-
rows.extend(convert(one))
289-
except Exception as e:
290-
rows.append(_FailedRow(e))
260+
261+
def convert_range(lo: int, hi: int, error: Optional[Exception]) -> None:
262+
# ``error`` is the failure already seen for [lo, hi), if it was tried.
263+
if error is None:
264+
piece = raw[lo:hi] if isinstance(raw, list) else raw.slice(lo, hi - lo)
265+
try:
266+
rows.extend(convert(piece))
267+
return
268+
except Exception as e:
269+
error = e
270+
if hi - lo == 1:
271+
rows.append(_FailedRow(error))
272+
return
273+
mid = (lo + hi) // 2
274+
convert_range(lo, mid, None)
275+
convert_range(mid, hi, None)
276+
277+
convert_range(0, n, batch_error)
291278
failed = sum(isinstance(r, _FailedRow) for r in rows)
292279
# One summary line, no traceback: the error itself still reaches the caller
293280
# when the failing row is fetched. The full traceback is only logged at DEBUG.
294281
# ``log_level`` lets the caller log this at WARNING only once per result set.
295282
logger.log(
296283
log_level,
297-
"fetchone(): failed to convert %d of %d row(s) in a batch%s; each failing "
284+
"fetchone(): failed to convert %d of %d row(s) in a batch; each failing "
298285
"row raises when fetched (%s: %s)",
299286
failed,
300287
n,
301-
(
302-
" (first %d rows all failed, rest not retried)"
303-
% _FETCHONE_FALLBACK_PROBE_ROWS
304-
if systematic
305-
else ""
306-
),
307288
type(batch_error).__name__,
308289
batch_error,
309290
)
@@ -379,7 +360,7 @@ def _fetchone_buffered(
379360
rows = convert(raw)
380361
except Exception as batch_error:
381362
# The queue has already moved past this batch, so it can't be
382-
# re-read. Fall back to converting it one row at a time so that,
363+
# re-read. Fall back to bisecting it down to single rows so that,
383364
# as before batching, only the row(s) that fail to convert are
384365
# lost: each failing row is parked as a _FailedRow and re-raises
385366
# its own error when it is reached.

‎tests/unit/test_fetches.py‎

Lines changed: 45 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -738,19 +738,16 @@ def test_fetchone_convert_failure_warns_once_per_result_set(self):
738738
[logging.WARNING, logging.DEBUG, logging.DEBUG],
739739
)
740740

741-
def test_fetchone_systematic_convert_failure_stops_per_row_retries(self):
741+
def test_fetchone_convert_failure_on_every_row_bisects_to_single_rows(self):
742742
from unittest.mock import patch
743-
from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS
744743

745-
batch_error = ValueError("batch failure")
744+
shared_error = ValueError("batch failure")
746745
calls = []
747746

748747
def convert(self_, table):
749748
calls.append(table.num_rows)
750-
if len(calls) == 1:
751-
raise batch_error
752-
# Same type and message as the batch error: a systematic failure.
753-
raise ValueError("batch failure")
749+
# The converter re-raises one shared exception instance for every row.
750+
raise shared_error
754751

755752
rs = self.make_columnar_result_set(50, arraysize=20)
756753
with patch.object(
@@ -763,27 +760,57 @@ def convert(self_, table):
763760
for i in range(20):
764761
with self.assertRaises(ValueError) as ctx:
765762
rs.fetchone()
766-
if i >= _FETCHONE_FALLBACK_PROBE_ROWS:
767-
self.assertIs(ctx.exception, batch_error)
768-
tb, depth = ctx.exception.__traceback__, 0
769-
while tb is not None:
770-
tb, depth = tb.tb_next, depth + 1
771-
tb_depths.append(depth)
763+
self.assertIs(ctx.exception, shared_error)
764+
tb, depth = ctx.exception.__traceback__, 0
765+
while tb is not None:
766+
tb, depth = tb.tb_next, depth + 1
767+
tb_depths.append(depth)
772768
self.assertEqual(rs.rownumber, i + 1)
773-
# Re-raising the shared batch error must not grow its traceback per row.
769+
# Re-raising a shared error must not grow its traceback per row.
774770
self.assertEqual(len(set(tb_depths)), 1)
775-
# One batch attempt plus only the probe rows were converted individually.
776-
self.assertEqual(len(calls), 1 + _FETCHONE_FALLBACK_PROBE_ROWS)
771+
# Bisection of an all-failing batch of 20 tries every node of the tree once:
772+
# 2 * 20 - 1 conversions (including the initial batch attempt).
773+
self.assertEqual(calls[0], 20)
774+
self.assertEqual(len(calls), 2 * 20 - 1)
775+
self.assertEqual(sum(1 for c in calls if c == 1), 20)
777776
self.assertTrue(any("20 of 20 row(s)" in m for m in logs.output))
778777
self.assertTrue(any("batch failure" in m for m in logs.output))
779778
self.assertEqual(rs.fetchone()[0], 20)
780779

780+
def test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows(self):
781+
from unittest.mock import patch
782+
783+
real_convert = ThriftResultSet._convert_columnar_table
784+
# Enough leading bad rows, all failing with the same message as the batch
785+
# (e.g. "year 0 is out of range"), to look like a systematic failure.
786+
n_bad = 8
787+
788+
def convert(self_, table):
789+
rows = real_convert(self_, table)
790+
if any(r[0] < n_bad for r in rows):
791+
raise ValueError("year 0 is out of range")
792+
return rows
793+
794+
rs = self.make_columnar_result_set(20, arraysize=20)
795+
with patch.object(
796+
ThriftResultSet,
797+
"_convert_columnar_table",
798+
autospec=True,
799+
side_effect=convert,
800+
), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
801+
for _ in range(n_bad):
802+
with self.assertRaises(ValueError):
803+
rs.fetchone()
804+
# Every row after the bad run still converts and is served.
805+
self.assertEqual([r[0] for r in rs.fetchall()], list(range(n_bad, 20)))
806+
self.assertTrue(any("{} of 20 row(s)".format(n_bad) in m for m in logs.output))
807+
self.assertEqual(rs.rownumber, 20)
808+
781809
def test_fetchone_leading_bad_rows_do_not_drop_later_good_rows(self):
782810
from unittest.mock import patch
783-
from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS
784811

785812
real_convert = ThriftResultSet._convert_columnar_table
786-
n_bad = _FETCHONE_FALLBACK_PROBE_ROWS + 1
813+
n_bad = 6
787814

788815
def convert(self_, table):
789816
rows = real_convert(self_, table)

0 commit comments

Comments
 (0)