Skip to content

Commit cfbc6f6

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

2 files changed

Lines changed: 64 additions & 3 deletions

File tree

‎src/databricks/sql/result_set.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,10 @@
4141
_FETCHONE_BUFFER_INITIAL_ROWS = 1
4242
_FETCHONE_BUFFER_GROWTH_FACTOR = 10
4343

44+
# Conversion errors that are not caused by row data. fetchone() re-raises these
45+
# straight away instead of bisecting the batch to find the failing row(s).
46+
_NON_DATA_CONVERSION_ERRORS = (MemoryError, RecursionError)
47+
4448

4549
class _FailedRow:
4650
"""Placeholder in the fetchone() row buffer for a row that failed to convert."""
@@ -379,6 +383,10 @@ def _fetchone_buffered(
379383
raw = fetch_raw(size)
380384
try:
381385
rows = convert(raw)
386+
except _NON_DATA_CONVERSION_ERRORS:
387+
# Resource / interpreter failures say nothing about individual
388+
# rows; retrying smaller slices would only repeat (or worsen) them.
389+
raise
382390
except Exception as batch_error:
383391
# The queue has already moved past this batch, so it can't be
384392
# re-read. Fall back to bisecting it down to single rows so that,
@@ -399,6 +407,12 @@ def _fetchone_buffered(
399407
raise
400408
self._fetchone_conversion_warned = True
401409
rows = fallback
410+
if all(isinstance(r, _FailedRow) for r in rows):
411+
# Every row failed: likely systemic (a column type or library
412+
# problem that hits each row). Restart the refill ramp so later
413+
# refills convert one row at a time, as before batching, rather
414+
# than paying a full ~2n-conversion bisection on every refill.
415+
self._fetchone_refill_rows = _FETCHONE_BUFFER_INITIAL_ROWS
402416
# Check for drift before the empty-result early return, so a converter
403417
# that yields [] for a non-empty chunk raises instead of silently
404418
# ending iteration and losing those rows.

‎tests/unit/test_fetches.py‎

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -725,19 +725,66 @@ def test_fetchone_convert_failure_warns_once_per_result_set(self):
725725
"_convert_columnar_table",
726726
side_effect=ValueError("bad value"),
727727
), self.assertLogs("databricks.sql.result_set", level="DEBUG") as logs:
728-
# Three refill batches, every row failing the same way.
728+
# One 4-row refill, then (ramp reset after an all-failing batch) eight
729+
# 1-row refills, every row failing the same way.
729730
for _ in range(12):
730731
with self.assertRaises(ValueError):
731732
rs.fetchone()
732733
summaries = [
733734
r for r in logs.records if "failed to convert" in r.getMessage()
734735
]
735-
self.assertEqual(len(summaries), 3)
736+
self.assertEqual(len(summaries), 9)
736737
self.assertEqual(
737738
[r.levelno for r in summaries],
738-
[logging.WARNING, logging.DEBUG, logging.DEBUG],
739+
[logging.WARNING] + [logging.DEBUG] * 8,
739740
)
740741

742+
def test_fetchone_all_failing_batch_resets_refill_ramp(self):
743+
from unittest.mock import patch
744+
745+
calls = []
746+
747+
def convert(self_, table):
748+
calls.append(table.num_rows)
749+
raise ValueError("systemic failure")
750+
751+
rs = self.make_columnar_result_set(30, arraysize=10)
752+
with patch.object(
753+
ThriftResultSet,
754+
"_convert_columnar_table",
755+
autospec=True,
756+
side_effect=convert,
757+
), self.assertLogs("databricks.sql.result_set", level="DEBUG"):
758+
for _ in range(15):
759+
with self.assertRaises(ValueError):
760+
rs.fetchone()
761+
# First full refill is bisected (2 * 10 - 1 conversions); after that every
762+
# refill is a single row costing exactly one conversion.
763+
self.assertEqual(len(calls), 2 * 10 - 1 + 5)
764+
self.assertEqual(calls[-5:], [1] * 5)
765+
self.assertEqual(rs.rownumber, 15)
766+
767+
def test_fetchone_memory_error_is_not_bisected(self):
768+
from unittest.mock import patch
769+
770+
calls = []
771+
772+
def convert(self_, table):
773+
calls.append(table.num_rows)
774+
raise MemoryError()
775+
776+
rs = self.make_columnar_result_set(10, arraysize=4)
777+
with patch.object(
778+
ThriftResultSet,
779+
"_convert_columnar_table",
780+
autospec=True,
781+
side_effect=convert,
782+
):
783+
with self.assertRaises(MemoryError):
784+
rs.fetchone()
785+
self.assertEqual(calls, [4])
786+
self.assertEqual(rs._row_buffer_count(), 0)
787+
741788
def test_fetchone_convert_failure_on_every_row_bisects_to_single_rows(self):
742789
from unittest.mock import patch
743790

0 commit comments

Comments
 (0)