diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py index 15f78bf5f..e5a32a37c 100644 --- a/src/databricks/sql/backend/kernel/result_set.py +++ b/src/databricks/sql/backend/kernel/result_set.py @@ -189,38 +189,77 @@ def _drain(self) -> pyarrow.Table: # ----- Arrow fetches ----- + def _fetchmany_arrow_unbuffered(self, size: int) -> pyarrow.Table: + """Take up to ``size`` rows from the batch FIFO, bypassing the + fetchone() row buffer.""" + if size == 0: + return pyarrow.Table.from_batches([], schema=self._schema) + self._ensure_buffered(size) + return self._take_buffered(size) + + def _fetch_available(self, size: int) -> pyarrow.Table: + """Take up to ``size`` rows already buffered, for fetchone() refills. + + Unlike ``_fetchmany_arrow_unbuffered`` this does not keep pulling + batches until ``size`` rows are collected: it only pulls from the + kernel while the batch FIFO is empty, so latency to the first row + matches the pre-buffering behaviour. + """ + self._ensure_buffered(1) + return self._take_buffered(min(size, self._buffered_rows())) + + @staticmethod + def _concat_buffered_arrow( + buffered: Optional[pyarrow.Table], rest: pyarrow.Table + ) -> pyarrow.Table: + """Prepend rows drained from the fetchone() buffer to ``rest``. + + Both tables are built against the result set's ``_schema``, so a plain + ``pyarrow.concat_tables`` suffices (no ColumnTable / schema + promotion handling as in ``ThriftResultSet._prepend_buffered``). + """ + if buffered is None: + return rest + if rest.num_rows == 0: + return buffered + return pyarrow.concat_tables([buffered, rest]) + def fetchall_arrow(self) -> pyarrow.Table: - return self._drain() + buffered = self._row_buffer_take_raw() + return self._concat_buffered_arrow(buffered, self._drain()) def fetchmany_arrow(self, size: int) -> pyarrow.Table: if size < 0: raise ValueError(f"fetchmany_arrow size must be >= 0, got {size}") - if size == 0: - return pyarrow.Table.from_batches([], schema=self._schema) - self._ensure_buffered(size) - return self._take_buffered(size) + buffered = self._row_buffer_take_raw(size) + remaining = size - (buffered.num_rows if buffered is not None else 0) + if buffered is not None and remaining == 0: + return buffered + return self._concat_buffered_arrow( + buffered, self._fetchmany_arrow_unbuffered(remaining) + ) # ----- Row fetches ----- def fetchone(self) -> Optional[Row]: - self._ensure_buffered(1) - if self._buffered_rows() == 0: - return None - table = self._take_buffered(1) - rows = self._convert_arrow_table(table) - return rows[0] if rows else None + # Rows are converted ``arraysize`` at a time and served from the + # base-class row buffer; converting one row per call is very slow. + return self._fetchone_buffered( + self._fetch_available, self._convert_arrow_table + ) def fetchmany(self, size: int) -> List[Row]: if size < 0: raise ValueError(f"fetchmany size must be >= 0, got {size}") - if size == 0: - return [] - self._ensure_buffered(size) - table = self._take_buffered(size) - return self._convert_arrow_table(table) + return self._fetch_rows_with_buffer( + size, + lambda n: self._convert_arrow_table(self._fetchmany_arrow_unbuffered(n)), + ) def fetchall(self) -> List[Row]: - return self._convert_arrow_table(self._drain()) + return self._fetch_rows_with_buffer( + None, lambda _: self._convert_arrow_table(self._drain()) + ) def close(self) -> None: """Close the underlying kernel handle and notify the backend. @@ -234,6 +273,10 @@ def close(self) -> None: connection-close would otherwise issue a kernel call into an already-disposed session. """ + # Drop rows parked by fetchone() (and any lost-rows poison) so + # post-close fetches return nothing, matching the cleared batch FIFO below. + self._row_buffer_clear() + self._lost_rows_error = None if self._kernel_handle is None: return if not self.connection.open: diff --git a/src/databricks/sql/backend/sea/result_set.py b/src/databricks/sql/backend/sea/result_set.py index 17838ed81..38e4727c2 100644 --- a/src/databricks/sql/backend/sea/result_set.py +++ b/src/databricks/sql/backend/sea/result_set.py @@ -19,6 +19,7 @@ from databricks.sql.backend.sea.queue import JsonQueue, SeaResultSetQueueFactory from databricks.sql.backend.types import ExecuteResponse from databricks.sql.result_set import ResultSet +from databricks.sql.utils import concat_table_chunks logger = logging.getLogger(__name__) @@ -160,10 +161,7 @@ def fetchmany_json(self, size: int) -> List[List[str]]: if size < 0: raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") - results = self.results.next_n_rows(size) - self._next_row_index += len(results) - - return results + return self._fetch_raw_with_buffer(size) def fetchall_json(self) -> List[List[str]]: """ @@ -173,11 +171,48 @@ def fetchall_json(self) -> List[List[str]]: Columnar table containing all remaining rows """ - results = self.results.remaining_rows() - self._next_row_index += len(results) + return self._fetch_raw_with_buffer(None) + + @staticmethod + def _raw_num_rows(raw: Any) -> int: + return len(raw) if isinstance(raw, list) else raw.num_rows + def _fetch_raw_from_queue(self, size: Optional[int]) -> Any: + """Fetch ``size`` (all if None) raw rows from the queue, bypassing the row buffer.""" + if size is None: + results = self.results.remaining_rows() + else: + results = self.results.next_n_rows(size) + self._next_row_index += self._raw_num_rows(results) return results + def _fetch_available_from_queue(self, size: int) -> Any: + """Fetch up to ``size`` rows the queue already holds, for Arrow fetchone() refills. + + On CloudFetch, the request is capped at what the current chunk still holds + (matching Thrift) via the shared ``ResultSet._cap_to_current_chunk``, so a + refill is limited to one chunk rather than stitching several links into one + batch. It does not avoid waiting on the next download: reading to the end of + a chunk still makes ``next_n_rows`` advance to the next link in that call. + Only called while the fetchone() row buffer is empty. + """ + return self._fetch_raw_from_queue(self._cap_to_current_chunk(size)) + + def _fetch_raw_with_buffer(self, size: Optional[int]) -> Any: + """Fetch ``size`` (all if None) raw rows, draining the fetchone() buffer first.""" + buffered = self._row_buffer_take_raw(size) + if buffered is None: + return self._fetch_raw_from_queue(size) + remaining = None if size is None else size - self._raw_num_rows(buffered) + if remaining == 0: + return buffered + rest = self._fetch_raw_from_queue(remaining) + if self._raw_num_rows(rest) == 0: + return buffered + if isinstance(buffered, list): + return buffered + list(rest) + return concat_table_chunks([buffered, rest]) + def fetchmany_arrow(self, size: int) -> "pyarrow.Table": """ Fetch the next set of rows as an Arrow table. @@ -196,12 +231,10 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table": if size < 0: raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") - results = self.results.next_n_rows(size) + results = self._fetch_raw_with_buffer(size) if isinstance(self.results, JsonQueue): results = self._convert_json_to_arrow_table(results) - self._next_row_index += results.num_rows - return results def fetchall_arrow(self) -> "pyarrow.Table": @@ -209,12 +242,10 @@ def fetchall_arrow(self) -> "pyarrow.Table": Fetch all remaining rows as an Arrow table. """ - results = self.results.remaining_rows() + results = self._fetch_raw_with_buffer(None) if isinstance(self.results, JsonQueue): results = self._convert_json_to_arrow_table(results) - self._next_row_index += results.num_rows - return results def fetchone(self) -> Optional[Row]: @@ -227,11 +258,10 @@ def fetchone(self) -> Optional[Row]: """ if isinstance(self.results, JsonQueue): - res = self._create_json_table(self.fetchmany_json(1)) - else: - res = self._convert_arrow_table(self.fetchmany_arrow(1)) - - return res[0] if res else None + return self._fetchone_buffered(self.fetchmany_json, self._create_json_table) + return self._fetchone_buffered( + self._fetch_available_from_queue, self._convert_arrow_table + ) def fetchmany(self, size: int) -> List[Row]: """ @@ -247,10 +277,15 @@ def fetchmany(self, size: int) -> List[Row]: ValueError: If size is negative """ + if size < 0: + raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") if isinstance(self.results, JsonQueue): - return self._create_json_table(self.fetchmany_json(size)) - else: - return self._convert_arrow_table(self.fetchmany_arrow(size)) + return self._fetch_rows_with_buffer( + size, lambda n: self._create_json_table(self.fetchmany_json(n)) + ) + return self._fetch_rows_with_buffer( + size, lambda n: self._convert_arrow_table(self.fetchmany_arrow(n)) + ) def fetchall(self) -> List[Row]: """ @@ -261,6 +296,9 @@ def fetchall(self) -> List[Row]: """ if isinstance(self.results, JsonQueue): - return self._create_json_table(self.fetchall_json()) - else: - return self._convert_arrow_table(self.fetchall_arrow()) + return self._fetch_rows_with_buffer( + None, lambda _: self._create_json_table(self.fetchall_json()) + ) + return self._fetch_rows_with_buffer( + None, lambda _: self._convert_arrow_table(self.fetchall_arrow()) + ) diff --git a/src/databricks/sql/backend/sea/utils/filters.py b/src/databricks/sql/backend/sea/utils/filters.py index f125c409a..662c8b243 100644 --- a/src/databricks/sql/backend/sea/utils/filters.py +++ b/src/databricks/sql/backend/sea/utils/filters.py @@ -174,8 +174,9 @@ def _filter_arrow_result_set( raise ValueError(f"Column index {column_index} is out of bounds") column_name = result_set.description[column_index][0] - # Get all remaining rows as Arrow table and filter it - arrow_table = result_set.results.remaining_rows() + # Get all remaining rows as Arrow table and filter it. Go through + # fetchall_arrow() so any rows buffered by fetchone() are included. + arrow_table = result_set.fetchall_arrow() filtered_table = ResultSetFilter._filter_arrow_table( arrow_table, column_name, allowed_values, case_sensitive ) @@ -220,8 +221,9 @@ def _filter_json_result_set( if column_index >= len(result_set.description): raise ValueError(f"Column index {column_index} is out of bounds") - # Extract rows - all_rows = result_set.results.remaining_rows() + # Extract rows. Go through fetchall_json() so any rows buffered by + # fetchone() are included. + all_rows = result_set.fetchall_json() # Convert allowed values if case-insensitive if not case_sensitive: diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index f24a07505..bf1761230 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import List, Optional, TYPE_CHECKING, Tuple +from typing import Any, Callable, List, Optional, TYPE_CHECKING, Tuple import logging import pandas @@ -16,8 +16,9 @@ from databricks.sql.client import Connection from databricks.sql.backend.databricks_client import DatabricksClient from databricks.sql.types import Row -from databricks.sql.exc import RequestError, CursorAlreadyClosedError +from databricks.sql.exc import InternalError, RequestError, CursorAlreadyClosedError from databricks.sql.utils import ( + CloudFetchQueue, ColumnTable, ColumnQueue, concat_table_chunks, @@ -27,6 +28,44 @@ logger = logging.getLogger(__name__) +# Records that may contain result data go to the non-propagating "unsafe" logger +# (configured the same way in thrift_backend.py), so enabling DEBUG on +# ``databricks.sql`` does not leak row values into logs. +unsafe_logger = logging.getLogger("databricks.sql.unsafe") +unsafe_logger.propagate = False +if not unsafe_logger.handlers: + unsafe_logger.addHandler(logging.NullHandler()) + +# Upper bound on how many rows a fetchone() refill pulls and converts at once. The +# refill uses ``arraysize`` but is capped here so a cursor with a large arraysize +# (default 100000) doesn't materialise a full page just to return one row. +_FETCHONE_BUFFER_MAX_ROWS = 1000 + +# fetchone() refills grow geometrically from this size (1 -> 10 -> 100 -> ...) up to +# min(arraysize, _FETCHONE_BUFFER_MAX_ROWS). The first refill converts a single row, +# so a lone fetchone() (e.g. SQLAlchemy ``.first()`` / ``.scalar()`` or an EXISTS +# probe) costs the same as before batching, while iteration loops quickly reach the +# full batch size. +_FETCHONE_BUFFER_INITIAL_ROWS = 1 +_FETCHONE_BUFFER_GROWTH_FACTOR = 10 + +# Conversion errors that are not caused by row data. fetchone() re-raises these +# straight away instead of bisecting the batch to find the failing row(s). +_NON_DATA_CONVERSION_ERRORS = (MemoryError, RecursionError) + + +class _FailedRow: + """Placeholder in the fetchone() row buffer for a row that failed to convert.""" + + __slots__ = ("error", "traceback") + + def __init__(self, error: Exception): + self.error = error + # Snapshot the conversion traceback. A converter may raise the same + # exception instance for many rows; restoring this before each raise keeps + # its __traceback__ from growing (and pinning frames) on every fetchone(). + self.traceback = error.__traceback__ + class ResultSet(ABC): """ @@ -85,6 +124,24 @@ def __init__( self._arrow_schema_bytes = arrow_schema_bytes # Affected-row count for DML; None for SELECT / unreported. self.num_modified_rows = num_modified_rows + # Small row buffer that serves fetchone()/iteration. Rows are fetched and + # converted ``arraysize`` at a time instead of one at a time (converting a + # single row costs a full pandas round-trip). ``_row_buffer_raw`` holds the + # un-converted chunk (Arrow table, ColumnTable or JSON row list) aligned with + # ``_row_buffer_rows`` so the Arrow/columnar fetch methods can drain it too. + # Every fetch method drains this buffer first, preserving row order. + self._row_buffer_raw: Any = None + self._row_buffer_rows: List[Row] = [] + self._row_buffer_pos = 0 + # Set once a fetchone() conversion failure has been logged at WARNING. + self._fetchone_conversion_warned = False + # Size of the next fetchone() refill; grows on each refill (see + # ``_FETCHONE_BUFFER_GROWTH_FACTOR``). + self._fetchone_refill_rows = _FETCHONE_BUFFER_INITIAL_ROWS + # Set when a fetchone() refill chunk was read from the queue but lost + # before it could be buffered (e.g. KeyboardInterrupt / MemoryError during + # conversion). Later fetches raise instead of silently skipping those rows. + self._lost_rows_error: Optional[str] = None def __iter__(self): while True: @@ -136,6 +193,323 @@ def _convert_arrow_table(self, table): def rownumber(self): return self._next_row_index + # ----- fetchone() row buffer helpers ----- + + def _check_no_lost_rows(self) -> None: + """Raise if an earlier fetchone() refill lost rows (see ``_lost_rows_error``).""" + if self._lost_rows_error is not None: + raise InternalError(self._lost_rows_error) + + def _row_buffer_count(self) -> int: + return len(self._row_buffer_rows) - self._row_buffer_pos + + def _row_buffer_clear(self) -> None: + self._row_buffer_raw = None + self._row_buffer_rows = [] + self._row_buffer_pos = 0 + + def _row_buffer_advance(self, n: Optional[int]) -> Tuple[int, int]: + """Advance the buffer by up to ``n`` rows (all if None); return (start, length).""" + available = self._row_buffer_count() + length = available if n is None else max(0, min(n, available)) + start = self._row_buffer_pos + self._row_buffer_pos += length + # Buffered rows only count as consumed once they are handed to the caller. + self._next_row_index += length + return start, length + + def _row_buffer_take_rows(self, n: Optional[int] = None) -> List[Row]: + """Pop up to ``n`` (all if None) converted rows from the buffer. + + If the requested range contains a ``_FailedRow``, only the rows up to and + including the first failed row are consumed and its error is raised. Rows + after it stay in the buffer for the next fetch. The good rows before it in + the same call are consumed with the failed fetch (a fetch either returns + all of its rows or raises), so ``rownumber`` stays aligned with what is + left in the buffer. + """ + self._check_no_lost_rows() + if self._row_buffer_count() == 0: + return [] + rows = self._row_buffer_rows + available = self._row_buffer_count() + length = available if n is None else max(0, min(n, available)) + start = self._row_buffer_pos + failed: Optional[_FailedRow] = None + for i in range(start, start + length): + row = rows[i] + if isinstance(row, _FailedRow): + failed = row + length = i - start + 1 + break + self._row_buffer_advance(length) + if self._row_buffer_count() == 0: + self._row_buffer_clear() + if failed is not None: + raise failed.error.with_traceback(failed.traceback) + return rows[start : start + length] + + def _fetch_rows_with_buffer( + self, + size: Optional[int], + fetch_rest: Callable[[Optional[int]], List[Row]], + ) -> List[Row]: + """Return ``size`` rows (all if None), draining the fetchone() buffer first. + + Every row-level fetchmany()/fetchall() goes through here, so rows the + fetchone() buffer already holds are always returned before newly fetched + ones. ``fetch_rest(n)`` must fetch up to ``n`` (all if None) rows + while bypassing the buffer. It is not called when the buffer already + covers ``size``. + """ + buffered = self._row_buffer_take_rows(size) + remaining = None if size is None else size - len(buffered) + if buffered and remaining == 0: + return buffered + rest = fetch_rest(remaining) + return buffered + rest if buffered else rest + + @staticmethod + def _raw_row_count(raw: Any) -> Optional[int]: + """Row count of a raw chunk (Arrow table, ColumnTable or list); None if unknown.""" + if raw is None: + return 0 + n = len(raw) if isinstance(raw, list) else getattr(raw, "num_rows", None) + return n if isinstance(n, int) else None + + @staticmethod + def _convert_rows_individually( + raw: Any, + convert: Callable[[Any], List[Row]], + batch_error: Exception, + log_level: int = logging.WARNING, + ) -> Optional[List[Any]]: + """Convert ``raw`` in smaller pieces after a batch conversion failed. + + The failed batch is bisected: each failing range is split in half and the + halves are converted separately, down to single rows. Only rows that fail + on their own are returned as ``_FailedRow`` placeholders (raised when that + row is handed out); every row that converts is kept, so a failure is never + attributed to a row that was not itself tried. Cost is about + ``2 * k * log2(n)`` conversions for ``k`` bad rows, and at most ``2n`` + when every row fails. Returns None if ``raw`` has no rows to split, in + which case the caller re-raises the batch error. + """ + n = ResultSet._raw_row_count(raw) + if not n: + return None + + rows: List[Any] = [] + + def convert_range(lo: int, hi: int, error: Optional[Exception]) -> None: + # ``error`` is the failure already seen for [lo, hi), if it was tried. + if error is None: + piece = raw[lo:hi] if isinstance(raw, list) else raw.slice(lo, hi - lo) + try: + rows.extend(convert(piece)) + return + except _NON_DATA_CONVERSION_ERRORS: + # Resource / interpreter failures are not row-specific: stop + # bisecting and propagate, as the top-level refill does. + raise + except Exception as e: + error = e + if hi - lo == 1: + rows.append(_FailedRow(error)) + return + mid = (lo + hi) // 2 + convert_range(lo, mid, None) + convert_range(mid, hi, None) + + convert_range(0, n, batch_error) + failed = sum(isinstance(r, _FailedRow) for r in rows) + # One summary line, no traceback: the error itself still reaches the caller + # when the failing row is fetched. Only the exception type is logged here, + # since conversion error messages often quote the offending result value; + # the message and full traceback only go to the non-propagating unsafe logger. + # ``log_level`` lets the caller log this at WARNING only once per result set. + logger.log( + log_level, + "fetchone(): failed to convert %d of %d row(s) in a batch; each failing " + "row raises when fetched (%s)", + failed, + n, + type(batch_error).__name__, + ) + unsafe_logger.debug( + "fetchone(): batch conversion error: %s", batch_error, exc_info=batch_error + ) + return rows + + def _cap_to_current_chunk(self, size: int) -> int: + """Cap ``size`` at the rows left in the current CloudFetch chunk, if any. + + Limits each fetchone() refill to a single CloudFetch chunk. It does not + avoid waiting on the next download: ``next_n_rows`` advances to the next + file as soon as the current one is fully read, inside the same call. + Shared by the Thrift and SEA fetchone() refills so the CloudFetchQueue + internals (``table`` / ``table_row_index``) are read in one place. + """ + if isinstance(self.results, CloudFetchQueue) and self.results.table: + available = self.results.table.num_rows - self.results.table_row_index + if available > 0: + return min(size, available) + return size + + def _row_buffer_take_raw(self, n: Optional[int] = None) -> Any: + """Pop up to ``n`` (all if None) un-converted rows from the buffer. + + Returns None when the buffer is empty, otherwise a slice of the raw chunk + (same type as what the queue returned: Arrow table, ColumnTable or list). + """ + self._check_no_lost_rows() + if self._row_buffer_count() == 0: + return None + raw = self._row_buffer_raw + start, length = self._row_buffer_advance(n) + if self._row_buffer_count() == 0: + self._row_buffer_clear() + if isinstance(raw, list): + return raw[start : start + length] + return raw.slice(start, length) + + def _fetchone_buffered( + self, fetch_raw: Callable[[int], Any], convert: Callable[[Any], List[Row]] + ) -> Optional[Row]: + """Serve fetchone() from the row buffer, refilling in small batches. + + Refill sizes start at ``_FETCHONE_BUFFER_INITIAL_ROWS`` and grow by + ``_FETCHONE_BUFFER_GROWTH_FACTOR`` on each refill, capped at ``arraysize`` + and ``_FETCHONE_BUFFER_MAX_ROWS``, so a single fetchone() converts only one + row while iteration settles into full batches. ``fetch_raw`` may return + fewer (Thrift and the + kernel backend only serve what is already buffered, refilling at most until + one row is available, so latency to the first row is unchanged). + + ``_next_row_index`` keeps counting rows handed to the caller: rows parked in + the buffer are subtracted on refill and added back as they are taken. The + queue is only read again once the buffer is empty, so at that point the + index equals the queue position again (Thrift relies on that for + ``expected_row_start_offset``). + + If conversion yields a different number of rows than the raw chunk held, + ``InternalError`` is raised and that whole chunk is dropped (the queue and + ``rownumber`` have already moved past it). The result set should be + considered unusable after that error. The same holds when conversion + raises one of ``_NON_DATA_CONVERSION_ERRORS`` (e.g. ``MemoryError``, + ``RecursionError``): the error is re-raised and the whole refill chunk + (up to ``_FETCHONE_BUFFER_MAX_ROWS`` rows, still counted in ``rownumber``) + is lost, so callers should not keep fetching from this result set. + The same applies to any ``BaseException`` (e.g. ``KeyboardInterrupt``) + raised during conversion. In all these cases every later fetch on this + result set raises ``InternalError`` rather than silently skipping the + lost rows. + """ + self._check_no_lost_rows() + if self._row_buffer_count() == 0: + cap = max(min(self.arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1) + size = max(min(self._fetchone_refill_rows, cap), 1) + raw = fetch_raw(size) + # Grow the ramp only once the fetch succeeded, so a transient fetch + # error (and a retried fetchone()) does not skip a ramp step. + self._fetchone_refill_rows = min( + size * _FETCHONE_BUFFER_GROWTH_FACTOR, _FETCHONE_BUFFER_MAX_ROWS + ) + # The queue (and rownumber) has already moved past ``raw``, but it is + # only buffered once conversion succeeds. Mark the result set as having + # lost rows until then, so any exception escaping below (including + # KeyboardInterrupt) makes later fetches fail loudly instead of + # resuming from the next chunk with a gap. + pending = self._raw_row_count(raw) + if pending != 0: + self._lost_rows_error = ( + "{} row(s) were lost because an earlier fetchone() failed while " + "converting them; this result set can no longer be read " + "consistently, re-execute the query".format( + "Some" if pending is None else pending + ) + ) + try: + rows = convert(raw) + except _NON_DATA_CONVERSION_ERRORS: + # Resource / interpreter failures say nothing about individual + # rows; retrying smaller slices would only repeat (or worsen) them. + # The queue has already moved past ``raw`` and it is not buffered, + # so the whole refill chunk is lost: as with the drift error below, + # the result set is unusable after this. + raise + except Exception as batch_error: + # The queue has already moved past this batch, so it can't be + # re-read. Fall back to bisecting it down to single rows so that, + # as before batching, only the row(s) that fail to convert are + # lost: each failing row is parked as a _FailedRow and re-raises + # its own error when it is reached. + # Why not just drop the batch and re-raise (as fetchmany does)? + # fetchone() used to convert one row at a time, so callers that + # catch a conversion error and keep iterating lost only that row. + # Dropping the whole refill would silently lose up to + # _FETCHONE_BUFFER_MAX_ROWS good rows per bad one, a behaviour + # change hidden behind a perf fix. Bisection keeps the old + # semantics; its ~2n-conversion worst case is bounded by the ramp + # reset below, and _FailedRow tracebacks are freed as they are + # consumed (or with the buffer). + # Warn once per result set; later failing batches (e.g. a column + # type that fails in every row) log at DEBUG to avoid a WARNING + # per refill for callers that catch and keep fetching. + fallback = self._convert_rows_individually( + raw, + convert, + batch_error, + log_level=( + logging.DEBUG + if self._fetchone_conversion_warned + else logging.WARNING + ), + ) + if fallback is None: + raise + self._fetchone_conversion_warned = True + rows = fallback + if all(isinstance(r, _FailedRow) for r in rows): + # Every row failed: likely systemic (a column type or library + # problem that hits each row). Restart the refill ramp so later + # refills convert one row at a time, as before batching, rather + # than paying a full ~2n-conversion bisection on every refill. + self._fetchone_refill_rows = _FETCHONE_BUFFER_INITIAL_ROWS + # Check for drift before the empty-result early return, so a converter + # that yields [] for a non-empty chunk raises instead of silently + # ending iteration and losing those rows. + raw_count = self._raw_row_count(raw) + if raw_count is None: + # Every real queue type (Arrow Table, ColumnTable, list) has a row + # count, so an unknown chunk type is itself a bug. Fail loudly + # rather than skip the drift check and risk rownumber misalignment. + raise InternalError( + "fetchone(): cannot determine the row count of a raw chunk of " + "type {}".format(type(raw).__name__) + ) + if raw_count != len(rows): + # fetch_raw advanced _next_row_index (and the Thrift fetch offset) + # by raw_count; the buffer must line up 1:1 with the raw chunk or + # rownumber and fetchmany_arrow slicing would silently drift. + # The raw chunk is not buffered, so its raw_count rows are lost + # (yet counted in rownumber): the result set is unusable after + # this error and later fetches raise (see ``_lost_rows_error``). + raise InternalError( + "fetchone(): converted {} row(s) from a chunk of {} raw row(s)".format( + len(rows), raw_count + ) + ) + if not rows: + return None + self._row_buffer_raw = raw + self._row_buffer_rows = rows + self._row_buffer_pos = 0 + self._next_row_index -= len(rows) + # ``raw`` is buffered now, so no rows were lost. + self._lost_rows_error = None + return self._row_buffer_take_rows(1)[0] + @property def is_staging_operation(self) -> bool: """Whether this result set represents a staging operation.""" @@ -189,6 +563,10 @@ def close(self) -> None: if isinstance(e.args[1], CursorAlreadyClosedError): logger.info("Operation was canceled by a prior request") finally: + # Drop rows parked by fetchone() (and any lost-rows poison) so + # post-close fetches return nothing. + self._row_buffer_clear() + self._lost_rows_error = None self.has_been_closed_server_side = True self.status = CommandState.CLOSED @@ -302,6 +680,14 @@ def _convert_columnar_table(self, table): return result + def _prepend_buffered(self, buffered, rest): + """Concatenate rows drained from the fetchone() buffer with freshly fetched ones.""" + if buffered is None: + return rest + if rest is None or rest.num_rows == 0: + return buffered + return concat_table_chunks([buffered, rest]) + def fetchmany_arrow(self, size: int) -> "pyarrow.Table": """ Fetch the next set of rows of a query result, returning a PyArrow table. @@ -309,7 +695,34 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table": An empty sequence is returned when no more rows are available. """ if size < 0: - raise ValueError("size argument for fetchmany is %s but must be >= 0", size) + raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") + buffered = self._row_buffer_take_raw(size) + remaining = size - (buffered.num_rows if buffered is not None else 0) + rest = self._fetchmany_arrow_from_queue(remaining) if remaining > 0 else None + if buffered is None and rest is None: + rest = self._fetchmany_arrow_from_queue(0) + return self._prepend_buffered(buffered, rest) + + def fetchmany_columnar(self, size: int): + """ + Fetch the next set of rows of a query result, returning a Columnar Table. + An empty sequence is returned when no more rows are available. + """ + if size < 0: + raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") + buffered = self._row_buffer_take_raw(size) + remaining = size - (buffered.num_rows if buffered is not None else 0) + rest = ( + self._fetchmany_columnar_from_queue(remaining) if remaining > 0 else None + ) + if buffered is None and rest is None: + rest = self._fetchmany_columnar_from_queue(0) + return self._prepend_buffered(buffered, rest) + + def _fetchmany_arrow_from_queue(self, size: int) -> "pyarrow.Table": + """Fetch up to ``size`` rows from the results queue, bypassing the row buffer.""" + if size < 0: + raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") # Hold 0-row chunks aside instead of appending them to ``partial_result_chunks``. # CloudFetchQueue may return a placeholder empty table whose schema does not @@ -343,13 +756,10 @@ def fetchmany_arrow(self, size: int) -> "pyarrow.Table": partial_result_chunks.append(zero_row_table) return concat_table_chunks(partial_result_chunks) - def fetchmany_columnar(self, size: int): - """ - Fetch the next set of rows of a query result, returning a Columnar Table. - An empty sequence is returned when no more rows are available. - """ + def _fetchmany_columnar_from_queue(self, size: int): + """Fetch up to ``size`` rows from the results queue, bypassing the row buffer.""" if size < 0: - raise ValueError("size argument for fetchmany is %s but must be >= 0", size) + raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") results = self.results.next_n_rows(size) n_remaining_rows = size - results.num_rows @@ -368,8 +778,51 @@ def fetchmany_columnar(self, size: int): return concat_table_chunks(partial_result_chunks) + def _fetch_available_from_queue(self, size: int): + """Fetch up to ``size`` rows the queue already holds, for fetchone() refills. + + Unlike ``_fetchmany_*_from_queue``, this does not keep calling the server + until ``size`` rows are collected: it only refills the queue while it is + empty, so fetchone() makes no more round-trips than it needs to return one + row (server pages can be smaller than ``size`` when they hit the byte limit). + On CloudFetch, the request is also capped at what the current chunk still + holds (see ``_cap_to_current_chunk``), so a refill never stitches rows from + several downloaded files into one batch. This only bounds the batch: reading + to the end of a chunk still makes ``next_n_rows`` advance to (and wait for) + the next file before it returns, exactly as the uncapped path did. + Only called while the fetchone() row buffer is empty. + """ + results = self.results.next_n_rows(self._cap_to_current_chunk(size)) + while ( + results.num_rows == 0 + and not self.has_been_closed_server_side + and self.has_more_rows + ): + self._fill_results_buffer() + results = self.results.next_n_rows(self._cap_to_current_chunk(size)) + self._next_row_index += results.num_rows + return results + def fetchall_arrow(self) -> "pyarrow.Table": """Fetch all (remaining) rows of a query result, returning them as a PyArrow table.""" + buffered = self._row_buffer_take_raw() + result_table = self._prepend_buffered( + buffered, self._fetchall_arrow_from_queue() + ) + # If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table + # Valid only for metadata commands result set + if isinstance(result_table, ColumnTable) and pyarrow: + data = { + name: col + for name, col in zip( + result_table.column_names, result_table.column_table + ) + } + return pyarrow.Table.from_pydict(data) + return result_table + + def _fetchall_arrow_from_queue(self): + """Drain the results queue (bypassing the row buffer) without type conversion.""" # Hold 0-row chunks aside instead of appending them to ``partial_result_chunks``. # CloudFetchQueue may return a placeholder empty table whose schema does not # match the real downloaded chunks; concatenating it would corrupt the result. @@ -394,21 +847,15 @@ def fetchall_arrow(self) -> "pyarrow.Table": if not partial_result_chunks: partial_result_chunks.append(zero_row_table) - result_table = concat_table_chunks(partial_result_chunks) - # If PyArrow is installed and we have a ColumnTable result, convert it to PyArrow Table - # Valid only for metadata commands result set - if isinstance(result_table, ColumnTable) and pyarrow: - data = { - name: col - for name, col in zip( - result_table.column_names, result_table.column_table - ) - } - return pyarrow.Table.from_pydict(data) - return result_table + return concat_table_chunks(partial_result_chunks) def fetchall_columnar(self): """Fetch all (remaining) rows of a query result, returning them as a Columnar table.""" + buffered = self._row_buffer_take_raw() + return self._prepend_buffered(buffered, self._fetchall_columnar_from_queue()) + + def _fetchall_columnar_from_queue(self): + """Drain the results queue as a Columnar table, bypassing the row buffer.""" results = self.results.remaining_rows() self._next_row_index += results.num_rows partial_result_chunks = [results] @@ -426,23 +873,24 @@ def fetchone(self) -> Optional[Row]: or None when no more data is available. """ if isinstance(self.results, ColumnQueue): - res = self._convert_columnar_table(self.fetchmany_columnar(1)) - else: - res = self._convert_arrow_table(self.fetchmany_arrow(1)) - - if len(res) > 0: - return res[0] - else: - return None + return self._fetchone_buffered( + self._fetch_available_from_queue, self._convert_columnar_table + ) + return self._fetchone_buffered( + self._fetch_available_from_queue, self._convert_arrow_table + ) def fetchall(self) -> List[Row]: """ Fetch all (remaining) rows of a query result, returning them as a list of rows. """ if isinstance(self.results, ColumnQueue): - return self._convert_columnar_table(self.fetchall_columnar()) - else: - return self._convert_arrow_table(self.fetchall_arrow()) + return self._fetch_rows_with_buffer( + None, lambda _: self._convert_columnar_table(self.fetchall_columnar()) + ) + return self._fetch_rows_with_buffer( + None, lambda _: self._convert_arrow_table(self.fetchall_arrow()) + ) def fetchmany(self, size: int) -> List[Row]: """ @@ -450,10 +898,15 @@ def fetchmany(self, size: int) -> List[Row]: An empty sequence is returned when no more rows are available. """ + if size < 0: + raise ValueError(f"size argument for fetchmany is {size} but must be >= 0") if isinstance(self.results, ColumnQueue): - return self._convert_columnar_table(self.fetchmany_columnar(size)) - else: - return self._convert_arrow_table(self.fetchmany_arrow(size)) + return self._fetch_rows_with_buffer( + size, lambda n: self._convert_columnar_table(self.fetchmany_columnar(n)) + ) + return self._fetch_rows_with_buffer( + size, lambda n: self._convert_arrow_table(self.fetchmany_arrow(n)) + ) @staticmethod def _get_schema_description(table_schema_message): diff --git a/tests/e2e/test_driver.py b/tests/e2e/test_driver.py index 319b7de5e..5c645af6c 100644 --- a/tests/e2e/test_driver.py +++ b/tests/e2e/test_driver.py @@ -942,6 +942,121 @@ def test_iterator_api(self, extra_params): for j in range(len(row)): assert row[j] == expected_results[i][j] + @contextmanager + def _count_row_conversions(self): + """Count every call that turns a fetched batch into ``Row`` objects + (Arrow, Thrift columnar and SEA JSON paths) — issue #981.""" + from unittest import mock + from databricks.sql.result_set import ResultSet, ThriftResultSet + from databricks.sql.backend.sea.result_set import SeaResultSet + + counter = {"calls": 0} + targets = [ + (ResultSet, "_convert_arrow_table"), + (ThriftResultSet, "_convert_columnar_table"), + (SeaResultSet, "_create_json_table"), + ] + patches = [] + for cls, name in targets: + orig = getattr(cls, name) + + def _wrapped(self_, *args, _orig=orig, **kwargs): + counter["calls"] += 1 + return _orig(self_, *args, **kwargs) + + patches.append(mock.patch.object(cls, name, _wrapped)) + for p in patches: + p.start() + try: + yield counter + finally: + for p in patches: + p.stop() + + _ROW_BUFFER_PARAMS = [ + {}, + { + "use_sea": True, + "use_cloud_fetch": False, + "enable_query_result_lz4_compression": False, + }, + { + "use_sea": True, + }, + ] + _ROW_BUFFER_N = 5000 + _ROW_BUFFER_QUERY = "SELECT id, CAST(id AS STRING) AS s FROM range({})" + + def _assert_rows_and_bound(self, rows, calls, arraysize): + n = self._ROW_BUFFER_N + assert len(rows) == n + assert [r[0] for r in rows] == list(range(n)) + assert all(r[1] == str(r[0]) for r in rows) + # Each refill converts up to min(arraysize, _FETCHONE_BUFFER_MAX_ROWS) + # rows, but refills are also capped at a server batch / CloudFetch chunk + # boundary, which can add one short refill per boundary. Allow 2x the + # ideal refill count plus a small constant for that slack — tight + # enough to catch tiny refills, not just one conversion per row (#981). + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + per_refill = max(min(arraysize, _FETCHONE_BUFFER_MAX_ROWS), 1) + ideal = -(-n // per_refill) + bound = 2 * ideal + 10 + assert calls <= bound, ( + "row conversion called {} times for {} rows (arraysize={}); " + "expected at most {} (issue #981)".format(calls, n, arraysize, bound) + ) + + @pytest.mark.parametrize("extra_params", _ROW_BUFFER_PARAMS) + def test_iteration_converts_rows_in_batches(self, extra_params): + with self.cursor(extra_params) as cursor: + cursor.execute(self._ROW_BUFFER_QUERY.format(self._ROW_BUFFER_N)) + with self._count_row_conversions() as counter: + rows = [row for row in cursor] + self._assert_rows_and_bound(rows, counter["calls"], cursor.arraysize) + + @pytest.mark.parametrize("extra_params", _ROW_BUFFER_PARAMS) + def test_fetchone_loop_converts_rows_in_batches(self, extra_params): + with self.cursor(extra_params) as cursor: + cursor.execute(self._ROW_BUFFER_QUERY.format(self._ROW_BUFFER_N)) + rows = [] + with self._count_row_conversions() as counter: + while True: + row = cursor.fetchone() + if row is None: + break + rows.append(row) + self._assert_rows_and_bound(rows, counter["calls"], cursor.arraysize) + + @pytest.mark.parametrize("extra_params", _ROW_BUFFER_PARAMS) + def test_fetchone_mixed_with_other_fetches_keeps_order(self, extra_params): + with self.cursor(extra_params) as cursor: + cursor.execute("SELECT * FROM range(10)") + assert cursor.fetchone()[0] == 0 + assert cursor.fetchone()[0] == 1 + assert cursor.rownumber == 2 + assert [r[0] for r in cursor.fetchmany(3)] == [2, 3, 4] + assert cursor.rownumber == 5 + assert cursor.fetchone()[0] == 5 + assert cursor.rownumber == 6 + table = cursor.fetchmany_arrow(2) + assert table.column(0).to_pylist() == [6, 7] + assert cursor.rownumber == 8 + assert [r[0] for r in cursor.fetchall()] == [8, 9] + assert cursor.rownumber == 10 + assert cursor.fetchone() is None + + @pytest.mark.parametrize("extra_params", _ROW_BUFFER_PARAMS) + def test_fetchone_then_fetchall_arrow_keeps_order(self, extra_params): + with self.cursor(extra_params) as cursor: + cursor.execute("SELECT * FROM range(10)") + assert cursor.fetchone()[0] == 0 + assert cursor.rownumber == 1 + table = cursor.fetchall_arrow() + assert table.column(0).to_pylist() == list(range(1, 10)) + assert cursor.rownumber == 10 + assert cursor.fetchone() is None + def test_temp_view_fetch(self): with self.cursor({}) as cursor: query = "create temporary view f as select * from range(10)" diff --git a/tests/e2e/test_kernel_backend.py b/tests/e2e/test_kernel_backend.py index ae0301fb6..e5132644d 100644 --- a/tests/e2e/test_kernel_backend.py +++ b/tests/e2e/test_kernel_backend.py @@ -183,6 +183,79 @@ def test_fetchall_arrow(conn): assert table.column_names == ["a", "b"] +def _count_arrow_conversions(): + """Patch ``ResultSet._convert_arrow_table`` to count calls (issue #981).""" + from unittest import mock + from databricks.sql.result_set import ResultSet + + orig = ResultSet._convert_arrow_table + counter = {"calls": 0} + + def _wrapped(self_, table): + counter["calls"] += 1 + return orig(self_, table) + + return counter, mock.patch.object(ResultSet, "_convert_arrow_table", _wrapped) + + +_ROW_BUFFER_N = 5000 +_ROW_BUFFER_ARRAYSIZE = 1000 + + +def _assert_rows_and_bound(rows, calls): + assert len(rows) == _ROW_BUFFER_N + assert [r[0] for r in rows] == list(range(_ROW_BUFFER_N)) + assert all(r[1] == str(r[0]) for r in rows) + # Each refill converts up to min(arraysize, _FETCHONE_BUFFER_MAX_ROWS) rows; + # server batch boundaries can add one short refill each. Allow 2x the ideal + # refill count plus a small constant — tight enough to catch tiny refills, + # not just one conversion per row (issue #981). + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + per_refill = max(min(_ROW_BUFFER_ARRAYSIZE, _FETCHONE_BUFFER_MAX_ROWS), 1) + ideal = -(-_ROW_BUFFER_N // per_refill) + bound = 2 * ideal + 10 + assert calls <= bound, ( + f"_convert_arrow_table called {calls} times for {_ROW_BUFFER_N} rows; " + f"expected at most {bound} (issue #981)" + ) + + +def test_iteration_converts_rows_in_batches(conn): + with conn.cursor(arraysize=_ROW_BUFFER_ARRAYSIZE) as cur: + cur.execute(f"SELECT id, CAST(id AS STRING) AS s FROM range({_ROW_BUFFER_N})") + counter, patch = _count_arrow_conversions() + with patch: + rows = [row for row in cur] + _assert_rows_and_bound(rows, counter["calls"]) + + +def test_fetchone_loop_converts_rows_in_batches(conn): + with conn.cursor(arraysize=_ROW_BUFFER_ARRAYSIZE) as cur: + cur.execute(f"SELECT id, CAST(id AS STRING) AS s FROM range({_ROW_BUFFER_N})") + counter, patch = _count_arrow_conversions() + rows = [] + with patch: + while True: + row = cur.fetchone() + if row is None: + break + rows.append(row) + _assert_rows_and_bound(rows, counter["calls"]) + + +def test_fetchone_mixed_with_other_fetches_keeps_order(conn): + with conn.cursor() as cur: + cur.execute("SELECT * FROM range(10)") + assert cur.fetchone()[0] == 0 + assert cur.fetchone()[0] == 1 + assert [r[0] for r in cur.fetchmany(3)] == [2, 3, 4] + assert cur.fetchone()[0] == 5 + assert cur.fetchmany_arrow(2).column(0).to_pylist() == [6, 7] + assert [r[0] for r in cur.fetchall()] == [8, 9] + assert cur.fetchone() is None + + @pytest.mark.parametrize("row_limit", [None, 0, 1, 5]) def test_cursor_row_limit(conn, row_limit): with conn.cursor(row_limit=row_limit) as cur: diff --git a/tests/unit/test_fetches.py b/tests/unit/test_fetches.py index e3a963355..bab2c2aa0 100644 --- a/tests/unit/test_fetches.py +++ b/tests/unit/test_fetches.py @@ -12,6 +12,7 @@ from databricks.sql.utils import ArrowQueue from databricks.sql.backend.thrift_backend import ThriftDatabricksClient from databricks.sql.result_set import ThriftResultSet +from databricks.sql.exc import InternalError class _StubArrowQueue: @@ -95,7 +96,8 @@ def make_dummy_result_set_from_initial_results(initial_results): return rs @staticmethod - def make_dummy_result_set_from_batch_list(batch_list): + def make_dummy_result_set_from_batch_list(batch_list, offsets=None): + """``offsets``, if given, records ``expected_row_start_offset`` per refill.""" batch_index = 0 def fetch_results( @@ -110,6 +112,8 @@ def fetch_results( chunk_id=0, ): nonlocal batch_index + if offsets is not None: + offsets.append(expected_row_start_offset) results = FetchTests.make_arrow_queue(batch_list[batch_index]) batch_index += 1 @@ -328,6 +332,129 @@ def test_fetchone_without_initial_results(self): dummy_result_set = self.make_dummy_result_set_from_batch_list(batch_list_2) self.assertEqual(dummy_result_set.fetchone(), None) + def test_mixed_fetches_across_pages_keep_order_and_refill_offsets(self): + # Issue #981: rows parked in the fetchone() buffer must not skew + # ``_next_row_index``, which Thrift sends as ``expected_row_start_offset`` + # on every refill. Pages are 3 rows; fetchone() refills 2 rows at a time, + # so buffer refills and server page boundaries interleave. + batch_list = [[[0], [1], [2]], [[3], [4], [5]], [[6], [7], [8]], [[9]]] + offsets = [] + rs = self.make_dummy_result_set_from_batch_list(batch_list, offsets=offsets) + rs.arraysize = 2 + self.assertEqual(offsets, [0]) + + self.assertSequenceEqual(rs.fetchone(), [0]) + self.assertEqual(rs.rownumber, 1) + + # Drains the buffered row, the rest of page 0, then refills page 1. + table = rs.fetchmany_arrow(3) + self.assertEqual(table.column(0).to_pylist(), [1, 2, 3]) + self.assertEqual(rs.rownumber, 4) + self.assertEqual(offsets, [0, 3]) + + self.assertSequenceEqual(rs.fetchone(), [4]) + self.assertSequenceEqual(rs.fetchone(), [5]) + self.assertEqual(rs.rownumber, 6) + self.assertEqual(offsets, [0, 3]) + + # Buffer empty and page 1 exhausted: the refill crosses into page 2. + self.assertSequenceEqual(rs.fetchone(), [6]) + self.assertEqual(rs.rownumber, 7) + self.assertEqual(offsets, [0, 3, 6]) + + # Drains the buffered row [7], the rest of page 2, then page 3. + self.assertEqualRowValues(rs.fetchall(), [[7], [8], [9]]) + self.assertEqual(rs.rownumber, 10) + self.assertEqual(offsets, [0, 3, 6, 9]) + + self.assertIsNone(rs.fetchone()) + self.assertEqual(rs.rownumber, 10) + self.assertEqual(offsets, [0, 3, 6, 9]) + + def test_close_drops_fetchone_row_buffer(self): + # Rows parked in the fetchone() buffer must not survive close(). The warm + # refill pulls every row into the buffer, so the queue itself is empty and + # any row returned after close() could only come from a stale buffer. + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + for fetch_after_close in ( + lambda rs: rs.fetchall(), + lambda rs: rs.fetchmany_arrow(10).to_pylist(), + ): + rs = self.make_dummy_result_set_from_initial_results([[1], [2], [3], [4]]) + rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS + self.assertSequenceEqual(rs.fetchone(), [1]) + self.assertEqual(rs._row_buffer_count(), 3) + rs.close() + self.assertEqual(rs._row_buffer_count(), 0) + self.assertEqual(fetch_after_close(rs), []) + self.assertIsNone(rs.fetchone()) + + def test_close_clears_lost_rows_error(self): + # A result set poisoned by a lost fetchone() refill must behave like any + # other closed result set: post-close fetches return nothing, not raise. + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + rs = self.make_dummy_result_set_from_initial_results([[1], [2], [3]]) + # Warm refill pulls every row into the buffer, so the queue is empty. + rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS + self.assertSequenceEqual(rs.fetchone(), [1]) + rs._lost_rows_error = "lost" + with self.assertRaises(InternalError): + rs.fetchone() + rs.close() + self.assertIsNone(rs._lost_rows_error) + self.assertIsNone(rs.fetchone()) + self.assertEqual(rs.fetchall(), []) + + def test_fetchone_refill_does_not_fetch_extra_pages(self): + # A refill must not make extra round-trips to fill ``arraysize`` when + # server pages are smaller (e.g. byte-limited): it serves what the current + # queue holds and only fetches a new page once the queue is empty. + batch_list = [[[0], [1]], [[2], [3]], [[4]]] + offsets = [] + rs = self.make_dummy_result_set_from_batch_list(batch_list, offsets=offsets) + rs.arraysize = 100 + self.assertEqual(offsets, [0]) + + self.assertSequenceEqual(rs.fetchone(), [0]) + self.assertEqual(offsets, [0]) + self.assertSequenceEqual(rs.fetchone(), [1]) + self.assertEqual(offsets, [0]) + self.assertSequenceEqual(rs.fetchone(), [2]) + self.assertEqual(offsets, [0, 2]) + self.assertEqual(rs.rownumber, 3) + self.assertEqualRowValues(rs.fetchall(), [[3], [4]]) + self.assertEqual(offsets, [0, 2, 4]) + self.assertIsNone(rs.fetchone()) + self.assertEqual(rs.rownumber, 5) + + def test_cloudfetch_fetchone_refill_capped_to_current_chunk(self): + # The fetchone() refill is capped at the rows left in the current CloudFetch + # chunk, so one refill never spans several files. (This only bounds the + # batch size; the mock queue does not model advancing to the next file.) + from databricks.sql.utils import ThriftCloudFetchQueue + + _, chunk = self.make_arrow_table([[i] for i in range(10)]) + queue = Mock(spec=ThriftCloudFetchQueue) + queue.table = chunk + queue.table_row_index = 8 + queue.next_n_rows.side_effect = lambda n: chunk.slice(8, n) + rs = self.make_dummy_result_set_from_queue_list( + [queue], + description=[("col0", "integer", None, None, None, None, None)], + ) + rs.arraysize = 100 + # Skip the adaptive refill ramp so the chunk cap, not the ramp, bounds it. + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS + + self.assertSequenceEqual(rs.fetchone(), [8]) + queue.next_n_rows.assert_called_once_with(2) + self.assertSequenceEqual(rs.fetchone(), [9]) + self.assertEqual(rs.rownumber, 2) + # Regression tests for fetchmany_arrow / fetchall_arrow handling of # the schemaless CloudFetch placeholder. def test_fetchall_arrow_drops_mismatched_empty_placeholder(self): @@ -391,5 +518,586 @@ def test_fetchmany_arrow_all_empty_returns_zero_row_table(self): self.assertEqual(result.num_rows, 0) +class ColumnarFetchoneBufferTests(unittest.TestCase): + """Issue #981: fetchone/iteration on the Thrift ColumnQueue path must convert + rows in batches (not one row at a time) while keeping order and rownumber.""" + + @staticmethod + def make_columnar_result_set(n_rows, arraysize, warm=True): + """Build a result set over ``n_rows`` columnar rows. + + With ``warm`` (the default) the adaptive fetchone() refill ramp is skipped, + so every refill is ``min(arraysize, _FETCHONE_BUFFER_MAX_ROWS)`` rows and the + batch-boundary assertions below stay simple. Pass ``warm=False`` to test + the ramp itself. + """ + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + from databricks.sql.utils import ColumnQueue, ColumnTable + + table = ColumnTable([list(range(n_rows))], ["col0"]) + mock_thrift_backend = Mock(spec=ThriftDatabricksClient) + mock_thrift_backend.fetch_results.return_value = (ColumnQueue(table), False, 0) + rs = ThriftResultSet( + connection=Mock(), + execute_response=ExecuteResponse( + command_id=None, + status=None, + has_been_closed_server_side=True, + description=[("col0", "integer", None, None, None, None, None)], + lz4_compressed=False, + is_staging_operation=False, + ), + thrift_client=mock_thrift_backend, + arraysize=arraysize, + ) + if warm: + rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS + return rs + + @staticmethod + def _count_conversions(): + from unittest.mock import patch + + return patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=ThriftResultSet._convert_columnar_table, + ) + + def test_iteration_converts_in_batches(self): + rs = self.make_columnar_result_set(100, arraysize=10) + with self._count_conversions() as m: + rows = [r[0] for r in rs] + self.assertEqual(rows, list(range(100))) + self.assertLessEqual(m.call_count, 100 // 10 + 2) + + def test_fetchone_loop_converts_in_batches(self): + rs = self.make_columnar_result_set(100, arraysize=10) + rows = [] + with self._count_conversions() as m: + while True: + row = rs.fetchone() + if row is None: + break + rows.append(row[0]) + self.assertEqual(rows, list(range(100))) + self.assertLessEqual(m.call_count, 100 // 10 + 2) + + def test_fetchone_refill_is_capped_for_large_arraysize(self): + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + n_rows = _FETCHONE_BUFFER_MAX_ROWS * 3 + rs = self.make_columnar_result_set(n_rows, arraysize=100000) + self.assertEqual(rs.fetchone()[0], 0) + self.assertEqual(rs._row_buffer_count(), _FETCHONE_BUFFER_MAX_ROWS - 1) + self.assertEqual(rs.rownumber, 1) + self.assertEqual([r[0] for r in rs.fetchall()], list(range(1, n_rows))) + self.assertEqual(rs.rownumber, n_rows) + + def test_first_fetchone_converts_a_single_row(self): + # A lone fetchone() (SQLAlchemy .first()/.scalar(), EXISTS probes) must not + # pull and convert a whole batch: the first refill is a single row. + rs = self.make_columnar_result_set(5000, arraysize=100000, warm=False) + with self._count_conversions() as m: + self.assertEqual(rs.fetchone()[0], 0) + self.assertEqual(m.call_count, 1) + self.assertEqual(m.call_args[0][1].num_rows, 1) + self.assertEqual(rs._row_buffer_count(), 0) + self.assertEqual(rs.rownumber, 1) + + def test_fetchone_refill_size_ramps_up_to_cap(self): + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + n_rows = 1 + 10 + 100 + _FETCHONE_BUFFER_MAX_ROWS * 2 + rs = self.make_columnar_result_set(n_rows, arraysize=100000, warm=False) + with self._count_conversions() as m: + rows = [r[0] for r in rs] + self.assertEqual(rows, list(range(n_rows))) + sizes = [c[0][1].num_rows for c in m.call_args_list] + self.assertEqual( + sizes[:5], + [1, 10, 100, _FETCHONE_BUFFER_MAX_ROWS, _FETCHONE_BUFFER_MAX_ROWS], + ) + self.assertEqual(rs.rownumber, n_rows) + + def test_fetchone_refill_ramp_is_capped_by_arraysize(self): + rs = self.make_columnar_result_set(10, arraysize=4, warm=False) + with self._count_conversions() as m: + rows = [r[0] for r in rs] + self.assertEqual(rows, list(range(10))) + sizes = [c[0][1].num_rows for c in m.call_args_list] + self.assertEqual(sizes[:4], [1, 4, 4, 1]) + + def test_mixed_fetches_keep_order_and_rownumber(self): + rs = self.make_columnar_result_set(10, arraysize=4) + self.assertEqual(rs.fetchone()[0], 0) + self.assertEqual(rs.fetchone()[0], 1) + self.assertEqual(rs.rownumber, 2) + self.assertEqual([r[0] for r in rs.fetchmany(3)], [2, 3, 4]) + self.assertEqual(rs.rownumber, 5) + self.assertEqual(rs.fetchone()[0], 5) + self.assertEqual(rs.rownumber, 6) + self.assertEqual([r[0] for r in rs.fetchall()], [6, 7, 8, 9]) + self.assertEqual(rs.rownumber, 10) + self.assertIsNone(rs.fetchone()) + + def test_fetchone_convert_failure_loses_only_the_bad_row(self): + from unittest.mock import patch + + real_convert = ThriftResultSet._convert_columnar_table + + def convert_failing_on_2(self_, table): + rows = real_convert(self_, table) + if any(r[0] == 2 for r in rows): + raise ValueError("bad value") + return rows + + rs = self.make_columnar_result_set(10, arraysize=4) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert_failing_on_2, + ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs: + self.assertEqual(rs.fetchone()[0], 0) + self.assertEqual(rs.fetchone()[0], 1) + with self.assertRaises(ValueError): + rs.fetchone() + # Only the failing row is lost; the rest of the batch is still served. + self.assertEqual(rs.rownumber, 3) + self.assertEqual(rs.fetchone()[0], 3) + self.assertEqual([r[0] for r in rs.fetchall()], list(range(4, 10))) + self.assertTrue(any("1 of 4 row(s)" in m for m in logs.output)) + self.assertEqual(rs.rownumber, 10) + + def _patch_convert_failing_on_2(self): + from unittest.mock import patch + + real_convert = ThriftResultSet._convert_columnar_table + + def convert_failing_on_2(self_, table): + rows = real_convert(self_, table) + if any(r[0] == 2 for r in rows): + raise ValueError("bad value") + return rows + + return patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert_failing_on_2, + ) + + def test_fetchmany_over_failed_buffered_row_keeps_rows_after_it(self): + rs = self.make_columnar_result_set(10, arraysize=4) + with self._patch_convert_failing_on_2(), self.assertLogs( + "databricks.sql.result_set", level="WARNING" + ): + # Buffer now holds rows 1, 2 (failed), 3. + self.assertEqual(rs.fetchone()[0], 0) + with self.assertRaises(ValueError): + rs.fetchmany(3) + # Consumed up to and including the failed row; row 3 stays buffered. + self.assertEqual(rs.rownumber, 3) + self.assertEqual(rs._row_buffer_count(), 1) + self.assertEqual([r[0] for r in rs.fetchmany(3)], [3, 4, 5]) + self.assertEqual([r[0] for r in rs.fetchall()], [6, 7, 8, 9]) + self.assertEqual(rs.rownumber, 10) + + def test_fetchall_over_failed_buffered_row_keeps_rows_after_it(self): + rs = self.make_columnar_result_set(10, arraysize=4) + with self._patch_convert_failing_on_2(), self.assertLogs( + "databricks.sql.result_set", level="WARNING" + ): + self.assertEqual(rs.fetchone()[0], 0) + with self.assertRaises(ValueError): + rs.fetchall() + self.assertEqual(rs.rownumber, 3) + self.assertEqual([r[0] for r in rs.fetchall()], list(range(3, 10))) + self.assertEqual(rs.rownumber, 10) + + def test_fetchone_convert_failure_on_every_row_reraises_per_row(self): + from unittest.mock import patch + + rs = self.make_columnar_result_set(6, arraysize=4) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + side_effect=ValueError("bad value"), + ): + for expected_rownumber in range(1, 5): + with self.assertRaises(ValueError): + rs.fetchone() + self.assertEqual(rs.rownumber, expected_rownumber) + self.assertEqual(rs._row_buffer_count(), 0) + self.assertEqual(rs.fetchone()[0], 4) + + def test_fetchone_convert_failure_warns_once_per_result_set(self): + import logging + from unittest.mock import patch + + rs = self.make_columnar_result_set(12, arraysize=4) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + side_effect=ValueError("bad value"), + ), self.assertLogs("databricks.sql.result_set", level="DEBUG") as logs: + # One 4-row refill, then (ramp reset after an all-failing batch) eight + # 1-row refills, every row failing the same way. + for _ in range(12): + with self.assertRaises(ValueError): + rs.fetchone() + summaries = [ + r for r in logs.records if "failed to convert" in r.getMessage() + ] + self.assertEqual(len(summaries), 9) + self.assertEqual( + [r.levelno for r in summaries], + [logging.WARNING] + [logging.DEBUG] * 8, + ) + # The error message may quote result data, so it must stay out of WARNING. + for r in logs.records: + if r.levelno >= logging.WARNING: + self.assertNotIn("bad value", r.getMessage()) + self.assertIn("ValueError", r.getMessage()) + + def test_fetchone_all_failing_batch_resets_refill_ramp(self): + from unittest.mock import patch + + calls = [] + + def convert(self_, table): + calls.append(table.num_rows) + raise ValueError("systemic failure") + + rs = self.make_columnar_result_set(30, arraysize=10) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert, + ), self.assertLogs("databricks.sql.result_set", level="DEBUG"): + for _ in range(15): + with self.assertRaises(ValueError): + rs.fetchone() + # First full refill is bisected (2 * 10 - 1 conversions); after that every + # refill is a single row costing exactly one conversion. + self.assertEqual(len(calls), 2 * 10 - 1 + 5) + self.assertEqual(calls[-5:], [1] * 5) + self.assertEqual(rs.rownumber, 15) + + def test_fetchone_memory_error_is_not_bisected(self): + from unittest.mock import patch + + calls = [] + + def convert(self_, table): + calls.append(table.num_rows) + raise MemoryError() + + rs = self.make_columnar_result_set(10, arraysize=4) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert, + ): + with self.assertRaises(MemoryError): + rs.fetchone() + self.assertEqual(calls, [4]) + self.assertEqual(rs._row_buffer_count(), 0) + # The lost refill chunk must not be skipped silently on later fetches. + with self.assertRaises(InternalError): + rs.fetchone() + + def test_fetchone_memory_error_during_bisection_stops_bisecting(self): + from unittest.mock import patch + + calls = [] + + def convert(self_, table): + calls.append(table.num_rows) + if len(calls) == 1: + raise ValueError("bad value") + raise MemoryError() + + rs = self.make_columnar_result_set(10, arraysize=4) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert, + ): + with self.assertRaises(MemoryError): + rs.fetchone() + # The full batch fails with a data error, the first half hits MemoryError: + # bisection stops there instead of converting more slices. + self.assertEqual(calls, [4, 2]) + self.assertEqual(rs._row_buffer_count(), 0) + # The lost refill chunk must not be skipped silently on later fetches. + with self.assertRaises(InternalError): + rs.fetchone() + + def test_fetchone_interrupted_conversion_poisons_later_fetches(self): + from unittest.mock import patch + + rs = self.make_columnar_result_set(10, arraysize=4) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + side_effect=KeyboardInterrupt(), + ): + with self.assertRaises(KeyboardInterrupt): + rs.fetchone() + # Conversion works again, but the interrupted chunk is gone: every fetch + # must raise instead of resuming from the next chunk with a gap. + for fetch in ( + rs.fetchone, + lambda: rs.fetchmany(2), + rs.fetchall, + lambda: rs.fetchmany_arrow(2), + rs.fetchall_arrow, + ): + with self.assertRaises(InternalError) as ctx: + fetch() + self.assertIn("4 row(s) were lost", str(ctx.exception)) + + def test_fetchone_successful_refill_does_not_poison(self): + rs = self.make_columnar_result_set(10, arraysize=4) + rows = [rs.fetchone() for _ in range(3)] + self.assertEqual([r[0] for r in rows], [0, 1, 2]) + self.assertIsNone(rs._lost_rows_error) + self.assertEqual([r[0] for r in rs.fetchall()], list(range(3, 10))) + + def test_fetchone_convert_failure_on_every_row_bisects_to_single_rows(self): + from unittest.mock import patch + + shared_error = ValueError("batch failure") + calls = [] + + def convert(self_, table): + calls.append(table.num_rows) + # The converter re-raises one shared exception instance for every row. + raise shared_error + + rs = self.make_columnar_result_set(50, arraysize=20) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert, + ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs: + tb_depths = [] + for i in range(20): + with self.assertRaises(ValueError) as ctx: + rs.fetchone() + self.assertIs(ctx.exception, shared_error) + tb, depth = ctx.exception.__traceback__, 0 + while tb is not None: + tb, depth = tb.tb_next, depth + 1 + tb_depths.append(depth) + self.assertEqual(rs.rownumber, i + 1) + # Re-raising a shared error must not grow its traceback per row. + self.assertEqual(len(set(tb_depths)), 1) + # Bisection of an all-failing batch of 20 tries every node of the tree once: + # 2 * 20 - 1 conversions (including the initial batch attempt). + self.assertEqual(calls[0], 20) + self.assertEqual(len(calls), 2 * 20 - 1) + self.assertEqual(sum(1 for c in calls if c == 1), 20) + self.assertTrue(any("20 of 20 row(s)" in m for m in logs.output)) + # Only the exception type is logged at WARNING; the message (which may quote + # result data) is kept out of it. + self.assertTrue(any("(ValueError)" in m for m in logs.output)) + self.assertFalse(any("batch failure" in m for m in logs.output)) + self.assertEqual(rs.fetchone()[0], 20) + + def test_fetchone_all_failing_max_batch_conversion_count_is_bounded(self): + """Document the worst-case cost of bisecting a full all-failing refill. + + When every row fails (a systemic converter / type problem), bisection tries + every node of the tree: exactly ``2n - 1`` conversions for an ``n``-row + refill, i.e. about 2x the one-conversion-per-row cost of fetchone() before + batching. Short-circuiting on "same exception type in both halves" is + deliberately not done, since it would drop good rows that share the error + type (see test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows). + """ + from unittest.mock import patch + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + n = _FETCHONE_BUFFER_MAX_ROWS + calls = [] + + def convert(self_, table): + calls.append(table.num_rows) + raise ValueError("systemic failure") + + rs = self.make_columnar_result_set(n, arraysize=n) + # Skip the adaptive refill ramp so the first refill is a full batch. + rs._fetchone_refill_rows = n + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert, + ), self.assertLogs("databricks.sql.result_set", level="WARNING"): + with self.assertRaises(ValueError): + rs.fetchone() + self.assertEqual(calls[0], n) + self.assertEqual(len(calls), 2 * n - 1) + # Total rows converted across all attempts is O(n log n), not O(n^2). + self.assertLessEqual(sum(calls), n * (n - 1).bit_length() + n) + self.assertEqual(rs._row_buffer_count(), n - 1) + + def test_fetchone_failed_fetch_does_not_advance_refill_ramp(self): + from unittest.mock import patch + + rs = self.make_columnar_result_set(100, arraysize=100, warm=False) + initial = rs._fetchone_refill_rows + real_next_n_rows = rs.results.next_n_rows + with patch.object( + rs.results, + "next_n_rows", + side_effect=[RuntimeError("transient"), real_next_n_rows], + ): + with self.assertRaises(RuntimeError): + rs.fetchone() + # The failed fetch must not grow the ramp, so a retry refills the same size. + self.assertEqual(rs._fetchone_refill_rows, initial) + sizes = [] + with patch.object( + rs.results, + "next_n_rows", + side_effect=lambda n: sizes.append(n) or real_next_n_rows(n), + ): + self.assertEqual(rs.fetchone()[0], 0) + self.assertEqual(sizes, [initial]) + + def test_fetchone_same_message_bad_rows_do_not_drop_later_good_rows(self): + from unittest.mock import patch + + real_convert = ThriftResultSet._convert_columnar_table + # Enough leading bad rows, all failing with the same message as the batch + # (e.g. "year 0 is out of range"), to look like a systematic failure. + n_bad = 8 + + def convert(self_, table): + rows = real_convert(self_, table) + if any(r[0] < n_bad for r in rows): + raise ValueError("year 0 is out of range") + return rows + + rs = self.make_columnar_result_set(20, arraysize=20) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert, + ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs: + for _ in range(n_bad): + with self.assertRaises(ValueError): + rs.fetchone() + # Every row after the bad run still converts and is served. + self.assertEqual([r[0] for r in rs.fetchall()], list(range(n_bad, 20))) + self.assertTrue(any("{} of 20 row(s)".format(n_bad) in m for m in logs.output)) + self.assertEqual(rs.rownumber, 20) + + def test_fetchone_leading_bad_rows_do_not_drop_later_good_rows(self): + from unittest.mock import patch + + real_convert = ThriftResultSet._convert_columnar_table + n_bad = 6 + + def convert(self_, table): + rows = real_convert(self_, table) + bad = [r[0] for r in rows if r[0] < n_bad] + if bad: + # Each bad row fails with its own message, unlike the batch error. + raise ValueError("bad value in row {}".format(bad[0])) + return rows + + rs = self.make_columnar_result_set(20, arraysize=20) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert, + ), self.assertLogs("databricks.sql.result_set", level="WARNING") as logs: + for i in range(n_bad): + with self.assertRaises(ValueError) as ctx: + rs.fetchone() + self.assertEqual(str(ctx.exception), "bad value in row {}".format(i)) + # Rows after the leading bad run still convert and are served. + self.assertEqual([r[0] for r in rs.fetchall()], list(range(n_bad, 20))) + self.assertTrue(any("{} of 20 row(s)".format(n_bad) in m for m in logs.output)) + self.assertFalse(any("rest not retried" in m for m in logs.output)) + self.assertEqual(rs.rownumber, 20) + + def test_fetchone_raw_converted_count_mismatch_fails_loudly(self): + from unittest.mock import patch + from databricks.sql.exc import InternalError + + real_convert = ThriftResultSet._convert_columnar_table + + def convert_dropping_a_row(self_, table): + return real_convert(self_, table)[:-1] + + rs = self.make_columnar_result_set(10, arraysize=4) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + side_effect=convert_dropping_a_row, + ): + with self.assertRaises(InternalError): + rs.fetchone() + + def test_fetchone_empty_conversion_of_nonempty_chunk_fails_loudly(self): + from unittest.mock import patch + from databricks.sql.exc import InternalError + + rs = self.make_columnar_result_set(10, arraysize=4) + with patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + return_value=[], + ): + with self.assertRaises(InternalError): + rs.fetchone() + + def test_fetchone_unknown_raw_chunk_type_fails_loudly(self): + from unittest.mock import patch + from databricks.sql.exc import InternalError + + class _NoRowCount: + pass + + rs = self.make_columnar_result_set(10, arraysize=4) + with patch.object( + ThriftResultSet, + "_fetch_available_from_queue", + autospec=True, + return_value=_NoRowCount(), + ), patch.object( + ThriftResultSet, + "_convert_columnar_table", + autospec=True, + return_value=[("x",)], + ): + with self.assertRaises(InternalError) as ctx: + rs.fetchone() + self.assertIn("_NoRowCount", str(ctx.exception)) + + @pytest.mark.skipif(pa is None, reason="PyArrow is not installed") + def test_fetchone_then_fetchall_arrow_keeps_order(self): + rs = self.make_columnar_result_set(10, arraysize=4) + self.assertEqual(rs.fetchone()[0], 0) + self.assertEqual(rs.rownumber, 1) + table = rs.fetchall_arrow() + self.assertEqual(table.column(0).to_pylist(), list(range(1, 10))) + self.assertEqual(rs.rownumber, 10) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_filters.py b/tests/unit/test_filters.py index 4efe51f3e..7036044c0 100644 --- a/tests/unit/test_filters.py +++ b/tests/unit/test_filters.py @@ -5,6 +5,8 @@ import unittest from unittest.mock import MagicMock, patch +import pytest + from databricks.sql.backend.sea.utils.filters import ResultSetFilter @@ -16,9 +18,9 @@ def setUp(self): # Create a mock SeaResultSet self.mock_sea_result_set = MagicMock() - # Set up the remaining_rows method on the results attribute + # Set up fetchall_json, which the filters use to read all remaining rows self.mock_sea_result_set.results = MagicMock() - self.mock_sea_result_set.results.remaining_rows.return_value = [ + self.mock_sea_result_set.fetchall_json.return_value = [ ["catalog1", "schema1", "table1", "owner1", "2023-01-01", "TABLE", ""], ["catalog1", "schema1", "table2", "owner1", "2023-01-01", "VIEW", ""], [ @@ -153,5 +155,121 @@ def test_filter_tables_by_type(self): ) +class TestResultSetFilterWithFetchoneBuffer(unittest.TestCase): + """Rows buffered by fetchone() on a real SeaResultSet must reach the filter.""" + + N_ROWS = 12 + + def _row_type(self, i): + return "TABLE" if i % 2 == 0 else "VIEW" + + def _make_result_set(self, queue, result_format): + from databricks.sql.backend.sea.models.base import ResultData, ResultManifest + from databricks.sql.backend.sea.result_set import SeaResultSet + from databricks.sql.backend.types import CommandId, CommandState + + connection = MagicMock() + connection.open = True + connection.disable_pandas = True + sea_client = MagicMock() + sea_client.max_download_threads = 10 + + execute_response = MagicMock() + execute_response.command_id = CommandId.from_sea_statement_id("stmt-1") + execute_response.status = CommandState.SUCCEEDED + execute_response.has_been_closed_server_side = False + execute_response.description = [ + ("table_name", "string", None, None, None, None, None), + ("table_type", "string", None, None, None, None, None), + ] + execute_response.is_staging_operation = False + execute_response.lz4_compressed = False + execute_response.arrow_schema_bytes = None + + manifest = ResultManifest( + format=result_format.value, + schema={}, + total_row_count=self.N_ROWS, + total_byte_count=-1, + total_chunk_count=-1, + ) + with patch( + "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue", + return_value=queue, + ): + return SeaResultSet( + connection=connection, + execute_response=execute_response, + sea_client=sea_client, + result_data=ResultData(data=None, external_links=None), + manifest=manifest, + buffer_size_bytes=1000, + arraysize=100, + ) + + def _consume_two_rows(self, rs): + # 1st fetchone() refills 1 row, 2nd refills 10: rows 2..10 stay buffered + # and row 11 is still in the queue. + self.assertEqual(rs.fetchone()[0], "t0") + self.assertEqual(rs.fetchone()[0], "t1") + self.assertGreater(rs._row_buffer_count(), 0) + + def _expected_names(self): + return [ + "t{}".format(i) + for i in range(2, self.N_ROWS) + if self._row_type(i) == "TABLE" + ] + + def test_json_filter_includes_rows_buffered_by_fetchone(self): + from databricks.sql.backend.sea.queue import JsonQueue + from databricks.sql.backend.sea.utils.constants import ResultFormat + + data = [["t{}".format(i), self._row_type(i)] for i in range(self.N_ROWS)] + rs = self._make_result_set(JsonQueue(data), ResultFormat.JSON_ARRAY) + self._consume_two_rows(rs) + + with patch( + "databricks.sql.backend.sea.result_set.SeaResultSet" + ) as mock_sea_result_set_class: + ResultSetFilter._filter_json_result_set( + rs, 1, ["TABLE"], case_sensitive=True + ) + result_data = mock_sea_result_set_class.call_args.kwargs["result_data"] + + self.assertEqual([r[0] for r in result_data.data], self._expected_names()) + self.assertEqual(rs.rownumber, self.N_ROWS) + + def test_arrow_filter_includes_rows_buffered_by_fetchone(self): + pyarrow = pytest.importorskip("pyarrow") + from databricks.sql.backend.sea.utils.constants import ResultFormat + from databricks.sql.utils import ArrowQueue + + table = pyarrow.Table.from_pydict( + { + "table_name": ["t{}".format(i) for i in range(self.N_ROWS)], + "table_type": [self._row_type(i) for i in range(self.N_ROWS)], + } + ) + rs = self._make_result_set( + ArrowQueue(table, self.N_ROWS), ResultFormat.ARROW_STREAM + ) + self._consume_two_rows(rs) + + with patch( + "databricks.sql.backend.sea.result_set.SeaResultSet" + ) as mock_sea_result_set_class: + ResultSetFilter._filter_arrow_result_set( + rs, 1, ["TABLE"], case_sensitive=True + ) + result_data = mock_sea_result_set_class.call_args.kwargs["result_data"] + + filtered = pyarrow.ipc.open_stream(result_data.attachment).read_all() + self.assertEqual( + filtered.column("table_name").to_pylist(), self._expected_names() + ) + self.assertEqual(rs.rownumber, self.N_ROWS) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py index cfa2d3f2e..ef2dcd6d7 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/tests/unit/test_kernel_result_set.py @@ -43,7 +43,7 @@ def close(self): self.closed = True -def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet: +def _make_rs(handle, *, disable_pandas=False, arraysize=100) -> KernelResultSet: # The base ResultSet __init__ takes a `connection` ref it never # actually dereferences during these buffer tests, so a Mock is # fine. @@ -55,7 +55,7 @@ def _make_rs(handle, *, disable_pandas=False) -> KernelResultSet: backend=backend, kernel_handle=handle, command_id=CommandId.from_sea_statement_id("smoke-test"), - arraysize=100, + arraysize=arraysize, buffer_size_bytes=1024, ) @@ -184,6 +184,84 @@ def test_fetchall_rows(int_schema): assert [r[0] for r in rows] == [1, 2, 3] +def _ten_rows_handle(int_schema): + # Batch boundaries (3/4/3) deliberately don't line up with arraysize=4 + # so fetchone() refills span kernel batches. + return _FakeKernelHandle( + int_schema, + [ + _batch(int_schema, [0, 1, 2]), + _batch(int_schema, [3, 4, 5, 6]), + _batch(int_schema, [7, 8, 9]), + ], + ) + + +def test_mixed_fetches_keep_order_and_rownumber(int_schema): + rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4) + assert rs.fetchone()[0] == 0 + assert rs.fetchone()[0] == 1 + assert rs.rownumber == 2 + assert [r[0] for r in rs.fetchmany(3)] == [2, 3, 4] + assert rs.rownumber == 5 + assert rs.fetchone()[0] == 5 + assert rs.rownumber == 6 + assert [r[0] for r in rs.fetchall()] == [6, 7, 8, 9] + assert rs.rownumber == 10 + assert rs.fetchone() is None + assert rs.rownumber == 10 + + +def test_fetchone_refill_pulls_at_most_one_batch(int_schema): + # First-row latency: a fetchone() refill must not loop fetch_next_batch() + # until arraysize rows are buffered. + handle = _ten_rows_handle(int_schema) + rs = _make_rs(handle, arraysize=100) + assert rs.fetchone()[0] == 0 + assert len(handle._batches) == 2 + assert rs.fetchone()[0] == 1 + assert rs.fetchone()[0] == 2 + assert len(handle._batches) == 2 + assert rs.fetchone()[0] == 3 + assert len(handle._batches) == 1 + + +def test_fetchone_then_fetchall_arrow_keeps_order(int_schema): + rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4) + assert rs.fetchone()[0] == 0 + assert rs.rownumber == 1 + table = rs.fetchall_arrow() + assert table.column(0).to_pylist() == list(range(1, 10)) + assert rs.rownumber == 10 + + +def test_fetchone_then_fetchmany_arrow_keeps_order(int_schema): + rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4) + assert rs.fetchone()[0] == 0 + # Served entirely from the fetchone() row buffer (2 rows left in it: the + # refill only took the first kernel batch). + assert rs.fetchmany_arrow(2).column(0).to_pylist() == [1, 2] + assert rs.rownumber == 3 + # Row buffer is empty; continues from the batch FIFO. + assert rs.fetchmany_arrow(3).column(0).to_pylist() == [3, 4, 5] + assert rs.rownumber == 6 + assert rs.fetchmany_arrow(0).num_rows == 0 + assert rs.rownumber == 6 + assert rs.fetchmany_arrow(10).column(0).to_pylist() == [6, 7, 8, 9] + assert rs.rownumber == 10 + + +def test_fetchone_then_fetchmany_rows_spans_buffer(int_schema): + rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4) + assert rs.fetchone()[0] == 0 + # 2 rows come from the row buffer, the rest from the batch FIFO. + assert [r[0] for r in rs.fetchmany(6)] == [1, 2, 3, 4, 5, 6] + assert rs.rownumber == 7 + assert [r[0] for r in rs.fetchmany(10)] == [7, 8, 9] + assert rs.rownumber == 10 + assert rs.fetchmany(1) == [] + + def test_fetchmany_negative_raises(int_schema): rs = _make_rs(_FakeKernelHandle(int_schema, [])) with pytest.raises(ValueError): @@ -201,6 +279,25 @@ def test_close_is_idempotent_and_calls_handle(int_schema): rs.close() # second call is a no-op (kernel handle is None) +def test_close_drops_fetchone_row_buffer(int_schema): + # Rows parked in the fetchone() buffer must not survive close(). + rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4) + assert rs.fetchone()[0] == 0 + rs.close() + assert rs.fetchone() is None + assert rs.fetchall() == [] + + +def test_close_clears_lost_rows_error(int_schema): + # A poisoned result set must return nothing after close(), not keep raising. + rs = _make_rs(_ten_rows_handle(int_schema), arraysize=4) + rs._lost_rows_error = "lost" + rs.close() + assert rs._lost_rows_error is None + assert rs.fetchone() is None + assert rs.fetchall() == [] + + def test_empty_stream(int_schema): rs = _make_rs(_FakeKernelHandle(int_schema, [])) assert rs.fetchone() is None diff --git a/tests/unit/test_sea_result_set.py b/tests/unit/test_sea_result_set.py index 1c3e3b5b4..db65c4a77 100644 --- a/tests/unit/test_sea_result_set.py +++ b/tests/unit/test_sea_result_set.py @@ -308,6 +308,29 @@ def test_close_when_connection_closed( assert result_set.has_been_closed_server_side is True assert result_set.status == CommandState.CLOSED + @pytest.mark.parametrize("fetch_after_close", ["fetchall", "fetchmany_arrow"]) + def test_close_drops_fetchone_row_buffer( + self, result_set_with_data, fetch_after_close + ): + """Rows parked in the fetchone() buffer must not survive close().""" + if fetch_after_close == "fetchmany_arrow" and pyarrow is None: + pytest.skip("PyArrow is not installed") + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + rs = result_set_with_data + # Warm refill pulls all 5 rows into the buffer, emptying the queue, so any + # row returned after close() could only come from a stale buffer. + rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS + assert rs.fetchone()[0] == "value1" + assert rs._row_buffer_count() == 4 + rs.close() + assert rs._row_buffer_count() == 0 + if fetch_after_close == "fetchall": + assert rs.fetchall() == [] + else: + assert rs.fetchmany_arrow(10).num_rows == 0 + assert rs.fetchone() is None + def test_convert_json_types(self, result_set_with_data, sample_data): """Test the _convert_json_types method.""" # Call _convert_json_types @@ -482,6 +505,223 @@ def test_iteration(self, result_set_with_data, sample_data): assert rows[0].col2 == 1 assert rows[0].col3 is True + def _make_json_result_set( + self, mock_connection, mock_sea_client, execute_response, n_rows, arraysize + ): + data = [["v{}".format(i), str(i), "true"] for i in range(n_rows)] + result_data = ResultData(data=data, external_links=None, row_count=n_rows) + with patch( + "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue", + return_value=JsonQueue(data), + ): + return SeaResultSet( + connection=mock_connection, + execute_response=execute_response, + sea_client=mock_sea_client, + result_data=result_data, + manifest=self._create_empty_manifest(ResultFormat.JSON_ARRAY), + buffer_size_bytes=1000, + arraysize=arraysize, + ) + + @pytest.mark.parametrize("mode", ["iterate", "fetchone"]) + def test_json_fetchone_converts_in_batches( + self, mock_connection, mock_sea_client, execute_response, mode + ): + """Issue #981: per-row fetchone must not convert one row at a time.""" + rs = self._make_json_result_set( + mock_connection, mock_sea_client, execute_response, 100, arraysize=10 + ) + with patch.object( + SeaResultSet, + "_create_json_table", + autospec=True, + side_effect=SeaResultSet._create_json_table, + ) as m: + if mode == "iterate": + rows = list(rs) + else: + rows = [] + while True: + row = rs.fetchone() + if row is None: + break + rows.append(row) + assert [r.col2 for r in rows] == list(range(100)) + assert m.call_count <= 100 // 10 + 2 + + def test_json_row_fetches_keep_order_and_rownumber( + self, mock_connection, mock_sea_client, execute_response + ): + """Mixed fetchone/fetchmany/fetchall ordering without any Arrow path.""" + rs = self._make_json_result_set( + mock_connection, mock_sea_client, execute_response, 10, arraysize=4 + ) + assert rs.fetchone().col2 == 0 + assert rs.fetchone().col2 == 1 + assert rs.rownumber == 2 + assert [r.col2 for r in rs.fetchmany(3)] == [2, 3, 4] + assert rs.rownumber == 5 + assert rs.fetchone().col2 == 5 + assert rs.rownumber == 6 + assert [r.col2 for r in rs.fetchall()] == [6, 7, 8, 9] + assert rs.rownumber == 10 + assert rs.fetchone() is None + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") + def test_json_mixed_fetches_keep_order_and_rownumber( + self, mock_connection, mock_sea_client, execute_response + ): + rs = self._make_json_result_set( + mock_connection, mock_sea_client, execute_response, 10, arraysize=4 + ) + assert rs.fetchone().col2 == 0 + assert rs.fetchone().col2 == 1 + assert rs.rownumber == 2 + assert [r.col2 for r in rs.fetchmany(3)] == [2, 3, 4] + assert rs.rownumber == 5 + assert rs.fetchone().col2 == 5 + assert rs.rownumber == 6 + table = rs.fetchmany_arrow(2) + assert table.column(1).to_pylist() == [6, 7] + assert rs.rownumber == 8 + assert [r.col2 for r in rs.fetchall()] == [8, 9] + assert rs.rownumber == 10 + assert rs.fetchone() is None + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") + def test_json_fetchone_then_fetchall_arrow_keeps_order( + self, mock_connection, mock_sea_client, execute_response + ): + rs = self._make_json_result_set( + mock_connection, mock_sea_client, execute_response, 10, arraysize=4 + ) + assert rs.fetchone().col2 == 0 + assert rs.rownumber == 1 + table = rs.fetchall_arrow() + assert table.column(1).to_pylist() == list(range(1, 10)) + assert rs.rownumber == 10 + + def _make_arrow_result_set( + self, mock_connection, mock_sea_client, execute_response, n_rows, arraysize + ): + from databricks.sql.utils import ArrowQueue + + mock_connection.disable_pandas = True + table = pyarrow.Table.from_pydict( + { + "col1": ["v{}".format(i) for i in range(n_rows)], + "col2": list(range(n_rows)), + "col3": [True] * n_rows, + } + ) + result_data = ResultData(data=None, external_links=None, row_count=n_rows) + with patch( + "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue", + return_value=ArrowQueue(table, n_rows), + ): + return SeaResultSet( + connection=mock_connection, + execute_response=execute_response, + sea_client=mock_sea_client, + result_data=result_data, + manifest=self._create_empty_manifest(ResultFormat.ARROW_STREAM), + buffer_size_bytes=1000, + arraysize=arraysize, + ) + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") + def test_arrow_mixed_fetches_keep_order_and_rownumber( + self, mock_connection, mock_sea_client, execute_response + ): + """Buffered Arrow slice joined with the queue's next chunk via concat.""" + rs = self._make_arrow_result_set( + mock_connection, mock_sea_client, execute_response, 10, arraysize=4 + ) + assert rs.fetchone().col2 == 0 + assert rs.rownumber == 1 + # 3 rows left in the buffer; ask for 5 so 2 must come from the queue. + table = rs.fetchmany_arrow(5) + assert isinstance(table, pyarrow.Table) + assert table.column(1).to_pylist() == [1, 2, 3, 4, 5] + assert rs.rownumber == 6 + assert rs.fetchone().col2 == 6 + assert rs.rownumber == 7 + # Fully served from the buffer (remaining == 0 path). + assert rs.fetchmany_arrow(1).column(1).to_pylist() == [7] + assert rs.rownumber == 8 + assert [r.col2 for r in rs.fetchall()] == [8, 9] + assert rs.rownumber == 10 + assert rs.fetchone() is None + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") + def test_arrow_buffer_with_exhausted_queue_returns_buffer( + self, mock_connection, mock_sea_client, execute_response + ): + """Queue returns a 0-row slice after the buffer drains: buffered rows only.""" + rs = self._make_arrow_result_set( + mock_connection, mock_sea_client, execute_response, 3, arraysize=10 + ) + assert rs.fetchone().col2 == 0 + table = rs.fetchmany_arrow(5) + assert table.column(1).to_pylist() == [1, 2] + assert rs.rownumber == 3 + assert rs.fetchall_arrow().num_rows == 0 + assert rs.rownumber == 3 + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") + def test_arrow_fetchone_then_fetchall_arrow_keeps_order( + self, mock_connection, mock_sea_client, execute_response + ): + rs = self._make_arrow_result_set( + mock_connection, mock_sea_client, execute_response, 10, arraysize=4 + ) + assert rs.fetchone().col2 == 0 + table = rs.fetchall_arrow() + assert table.column(1).to_pylist() == list(range(1, 10)) + assert rs.rownumber == 10 + + @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed") + def test_cloudfetch_fetchone_refill_capped_to_current_chunk( + self, mock_connection, mock_sea_client, execute_response + ): + """fetchone() refill is capped at the rows left in the current chunk (batch size only).""" + from databricks.sql.backend.sea.queue import SeaCloudFetchQueue + + chunk = pyarrow.Table.from_pydict( + { + "col1": ["v{}".format(i) for i in range(10)], + "col2": list(range(10)), + "col3": [True] * 10, + } + ) + queue = Mock(spec=SeaCloudFetchQueue) + queue.table = chunk + queue.table_row_index = 8 + queue.next_n_rows.side_effect = lambda n: chunk.slice(8, n) + result_data = ResultData(data=None, external_links=[], row_count=10) + with patch( + "databricks.sql.backend.sea.queue.SeaResultSetQueueFactory.build_queue", + return_value=queue, + ): + rs = SeaResultSet( + connection=mock_connection, + execute_response=execute_response, + sea_client=mock_sea_client, + result_data=result_data, + manifest=self._create_empty_manifest(ResultFormat.ARROW_STREAM), + buffer_size_bytes=1000, + arraysize=100, + ) + # Skip the adaptive refill ramp so the chunk cap, not the ramp, bounds it. + from databricks.sql.result_set import _FETCHONE_BUFFER_MAX_ROWS + + rs._fetchone_refill_rows = _FETCHONE_BUFFER_MAX_ROWS + assert rs.fetchone().col2 == 8 + queue.next_n_rows.assert_called_once_with(2) + assert rs.fetchone().col2 == 9 + assert rs.rownumber == 2 + def test_is_staging_operation( self, mock_connection, mock_sea_client, execute_response ):