Skip to content

Commit 5e24726

Browse files
ai: apply changes for #1007 (3 review threads)
Addresses: - #4418954338 at src/databricks/sql/result_set.py:228 - #4418954347 at src/databricks/sql/result_set.py:203 - #4418954351 at src/databricks/sql/result_set.py:316 Signed-off-by: peco-engineer-bot[bot] <peco-engineer-bot[bot]@users.noreply.github.com>
1 parent 8e03ea6 commit 5e24726

2 files changed

Lines changed: 73 additions & 14 deletions

File tree

‎src/databricks/sql/result_set.py‎

Lines changed: 32 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -35,19 +35,25 @@
3535
_FETCHONE_BUFFER_MAX_ROWS = 1000
3636

3737
# 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.
38+
# If this many leading rows all fail too, with the same error type and message as
39+
# the batch, the failure is treated as systematic: the rest of the batch is not
40+
# retried row by row (each would cost a separate conversion) and those rows re-raise
41+
# the original batch error instead. Leading rows failing for other reasons do not
42+
# trigger the shortcut, so later good rows are still converted.
4143
_FETCHONE_FALLBACK_PROBE_ROWS = 5
4244

4345

4446
class _FailedRow:
4547
"""Placeholder in the fetchone() row buffer for a row that failed to convert."""
4648

47-
__slots__ = ("error",)
49+
__slots__ = ("error", "traceback")
4850

4951
def __init__(self, error: Exception):
5052
self.error = error
53+
# Snapshot the conversion traceback. Systematic failures park the same
54+
# exception instance in many rows; restoring this before each raise keeps
55+
# its __traceback__ from growing (and pinning frames) on every fetchone().
56+
self.traceback = error.__traceback__
5157

5258

5359
class ResultSet(ABC):
@@ -200,7 +206,7 @@ def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]:
200206
if isinstance(row, _FailedRow):
201207
# The taken rows are consumed either way, matching how a failed
202208
# conversion behaved before fetchone() buffering.
203-
raise row.error
209+
raise row.error.with_traceback(row.traceback)
204210
return taken
205211

206212
@staticmethod
@@ -219,20 +225,29 @@ def _convert_rows_individually(
219225
220226
Rows that fail are returned as ``_FailedRow`` placeholders so the error is
221227
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
228+
``_FETCHONE_FALLBACK_PROBE_ROWS`` rows all fail with the same error as the
229+
batch (same type and message), the failure is treated as systematic and the
230+
remaining rows are parked with ``batch_error`` instead of being converted one
231+
by one. Otherwise every row is tried individually. Returns None if ``raw`` has no rows to split, in
225232
which case the caller re-raises the batch error.
226233
"""
227234
n = ResultSet._raw_row_count(raw)
228235
if not n:
229236
return None
237+
def same_as_batch(r: Any) -> bool:
238+
return (
239+
isinstance(r, _FailedRow)
240+
and type(r.error) is type(batch_error)
241+
and str(r.error) == str(batch_error)
242+
)
243+
230244
rows: List[Any] = []
231245
systematic = False
232246
for i in range(n):
233-
if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(
234-
isinstance(r, _FailedRow) for r in rows
235-
):
247+
# Only skip the remaining per-row retries when every probe row failed
248+
# with the same error as the batch; leading rows that fail for their
249+
# own reasons must not cause later good rows to be dropped.
250+
if i == _FETCHONE_FALLBACK_PROBE_ROWS and all(same_as_batch(r) for r in rows):
236251
systematic = True
237252
rows.extend(_FailedRow(batch_error) for _ in range(n - i))
238253
break
@@ -242,16 +257,20 @@ def _convert_rows_individually(
242257
except Exception as e:
243258
rows.append(_FailedRow(e))
244259
failed = sum(isinstance(r, _FailedRow) for r in rows)
260+
# One summary line, no traceback: the error itself still reaches the caller
261+
# when the failing row is fetched. The full traceback is only logged at DEBUG.
245262
logger.warning(
246263
"fetchone(): failed to convert %d of %d row(s) in a batch%s; each failing "
247-
"row raises when fetched",
264+
"row raises when fetched (%s: %s)",
248265
failed,
249266
n,
250267
" (first %d rows all failed, rest not retried)" % _FETCHONE_FALLBACK_PROBE_ROWS
251268
if systematic
252269
else "",
253-
exc_info=batch_error,
270+
type(batch_error).__name__,
271+
batch_error,
254272
)
273+
logger.debug("fetchone(): batch conversion error", exc_info=batch_error)
255274
return rows
256275

257276
def _cap_to_current_chunk(self, size: int) -> int:

‎tests/unit/test_fetches.py‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -612,7 +612,8 @@ def convert(self_, table):
612612
calls.append(table.num_rows)
613613
if len(calls) == 1:
614614
raise batch_error
615-
raise ValueError("row failure")
615+
# Same type and message as the batch error: a systematic failure.
616+
raise ValueError("batch failure")
616617

617618
rs = self.make_columnar_result_set(50, arraysize=20)
618619
with patch.object(
@@ -621,18 +622,57 @@ def convert(self_, table):
621622
autospec=True,
622623
side_effect=convert,
623624
), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
625+
tb_depths = []
624626
for i in range(20):
625627
with self.assertRaises(ValueError) as ctx:
626628
rs.fetchone()
627629
if i >= _FETCHONE_FALLBACK_PROBE_ROWS:
628630
self.assertIs(ctx.exception, batch_error)
631+
tb, depth = ctx.exception.__traceback__, 0
632+
while tb is not None:
633+
tb, depth = tb.tb_next, depth + 1
634+
tb_depths.append(depth)
629635
self.assertEqual(rs.rownumber, i + 1)
636+
# Re-raising the shared batch error must not grow its traceback per row.
637+
self.assertEqual(len(set(tb_depths)), 1)
630638
# One batch attempt plus only the probe rows were converted individually.
631639
self.assertEqual(len(calls), 1 + _FETCHONE_FALLBACK_PROBE_ROWS)
632640
self.assertTrue(any("20 of 20 row(s)" in m for m in logs.output))
633641
self.assertTrue(any("batch failure" in m for m in logs.output))
634642
self.assertEqual(rs.fetchone()[0], 20)
635643

644+
def test_fetchone_leading_bad_rows_do_not_drop_later_good_rows(self):
645+
from unittest.mock import patch
646+
from databricks.sql.result_set import _FETCHONE_FALLBACK_PROBE_ROWS
647+
648+
real_convert = ThriftResultSet._convert_columnar_table
649+
n_bad = _FETCHONE_FALLBACK_PROBE_ROWS + 1
650+
651+
def convert(self_, table):
652+
rows = real_convert(self_, table)
653+
bad = [r[0] for r in rows if r[0] < n_bad]
654+
if bad:
655+
# Each bad row fails with its own message, unlike the batch error.
656+
raise ValueError("bad value in row {}".format(bad[0]))
657+
return rows
658+
659+
rs = self.make_columnar_result_set(20, arraysize=20)
660+
with patch.object(
661+
ThriftResultSet,
662+
"_convert_columnar_table",
663+
autospec=True,
664+
side_effect=convert,
665+
), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs:
666+
for i in range(n_bad):
667+
with self.assertRaises(ValueError) as ctx:
668+
rs.fetchone()
669+
self.assertEqual(str(ctx.exception), "bad value in row {}".format(i))
670+
# Rows after the leading bad run still convert and are served.
671+
self.assertEqual([r[0] for r in rs.fetchall()], list(range(n_bad, 20)))
672+
self.assertTrue(any("{} of 20 row(s)".format(n_bad) in m for m in logs.output))
673+
self.assertFalse(any("rest not retried" in m for m in logs.output))
674+
self.assertEqual(rs.rownumber, 20)
675+
636676
def test_fetchone_raw_converted_count_mismatch_fails_loudly(self):
637677
from unittest.mock import patch
638678
from databricks.sql.exc import InternalError

0 commit comments

Comments
 (0)