diff --git a/.github/workflows/code-quality-checks.yml b/.github/workflows/code-quality-checks.yml index 9eb60a02a..36532be55 100644 --- a/.github/workflows/code-quality-checks.yml +++ b/.github/workflows/code-quality-checks.yml @@ -202,6 +202,9 @@ jobs: - name: Drive use_kernel=True through the REAL wheel — routing (offline) if: steps.deps.outputs.venv-usable == 'true' run: .venv/bin/python -m pytest tests/unit/test_session.py -m realkernel -v + - name: Real kernel row values match the pandas path (offline) + if: steps.deps.outputs.venv-usable == 'true' + run: .venv/bin/python -m pytest tests/unit/test_kernel_result_set.py -m realkernel -v # --- Fallback path: same-repo PRs / cache miss use JFrog (unchanged) -- - name: Setup Poetry @@ -239,6 +242,10 @@ jobs: # instantiates the real KernelDatabricksClient (not a stub, not a # Thrift fallback). Fails loudly if the real wheel is shadowed. run: poetry run python -m pytest tests/unit/test_session.py -m realkernel -v + - name: Real kernel row values match the pandas path + if: steps.deps.outputs.cache-hit != 'true' + # Skipped until the installed kernel has fetch_all_rows. + run: poetry run python -m pytest tests/unit/test_kernel_result_set.py -m realkernel -v check-linting: runs-on: diff --git a/CHANGELOG.md b/CHANGELOG.md index 867be71ec..6690d2de2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - Speed up Thrift response handling: read replies in 1 MiB chunks and decode them with Thrift's C-accelerated protocol, and stop stringifying every request/response for the default no-op debug logger. Wide result schemas, metadata calls with many rows, and inline Arrow results benefit most. - Speed up fetching small results on the kernel backend (`use_kernel=True`): they are converted to rows without building a pandas DataFrame when that is cheaper and gives the same values, which cut client CPU per query by 47% in a local benchmark of 200 concurrent sessions running one-row queries against a mock server. - Cut more per-query work on the kernel backend: reuse the cursor description while a connection's statements return the same result schema, and build small results' rows straight from the record batches instead of first combining them into a table. +- With a kernel that provides `fetch_all_rows`, the kernel backend gets small results' rows as Python values built by the kernel and reuses the cursor description without converting the schema to pyarrow, which raised one-process `SELECT 1` throughput with 800 sessions by about 65% against a Real-Time warehouse. # 4.9.0 (2026-11-04) - Upgrade Databricks SQL Kernel to 1.4.0, which keeps Lakehouse Real-Time session affinity and flushes shared telemetry only after the final session closes. diff --git a/src/databricks/sql/backend/kernel/result_set.py b/src/databricks/sql/backend/kernel/result_set.py index c008f6eb3..9d0a3e58e 100644 --- a/src/databricks/sql/backend/kernel/result_set.py +++ b/src/databricks/sql/backend/kernel/result_set.py @@ -9,14 +9,17 @@ ``ExecutedAsyncStatement.await_result()``. No statement id; no cancel. -Both implement the same three methods this class actually calls: -``arrow_schema() / fetch_next_batch() / close()``. ``KernelResultSet`` -takes either via the ``kernel_handle`` parameter and treats them -uniformly — the connector's ``ResultSet`` contract doesn't need to -distinguish them. - -Buffer shape mirrors the prior ADBC POC's ``AdbcResultSet``: a FIFO -of pyarrow ``RecordBatch``es, fed one batch at a time from the +Both implement ``arrow_schema() / fetch_next_batch() / close()``. +``KernelResultSet`` takes either via the ``kernel_handle`` parameter and +treats them uniformly — the connector's ``ResultSet`` contract doesn't +need to distinguish them. Newer kernels also have ``schema_key()`` and +``fetch_all_rows()``: the schema key looks the cached description up +without calling ``arrow_schema()``, and the first row fetch (unless an +Arrow fetch came first) gets a small result's rows as Python values in +one ``fetch_all_rows()`` call. + +Otherwise the buffer mirrors the prior ADBC POC's ``AdbcResultSet``: a +FIFO of pyarrow ``RecordBatch``es, fed one batch at a time from the kernel as the connector calls ``fetch*``. ``fetchmany(n)`` slices within a batch when ``n`` is smaller than the kernel's natural batch size; ``fetchall`` drains the whole stream. @@ -124,14 +127,17 @@ class _ResultSchema: __slots__ = ( "schema", + "key", "description", "row_factory", "max_rows_without_pandas", "float_columns", ) - def __init__(self, schema: pyarrow.Schema): + def __init__(self, schema: pyarrow.Schema, key: Any = None): self.schema = schema + # The kernel's SchemaKey for ``schema``, if the kernel has them. + self.key = key self.description = tuple(description_from_arrow_schema(schema)) self.row_factory = Row(*[c[0] for c in self.description]) types = schema.types @@ -141,6 +147,20 @@ def __init__(self, schema: pyarrow.Schema): ) +def _schema_key(kernel_handle: Any) -> Any: + """The handle's ``SchemaKey``, or None for kernels without them (or if + getting it fails: it only saves work).""" + try: + schema_key = kernel_handle.schema_key + except AttributeError: + return None + try: + return schema_key() + except Exception: + logger.debug("schema_key() failed; using arrow_schema()", exc_info=True) + return None + + class _SchemaCache: """Keeps the last ``_ResultSchema`` of a connection, since repeated statements return equal schemas.""" @@ -150,11 +170,22 @@ class _SchemaCache: def __init__(self) -> None: self._last: Optional[_ResultSchema] = None - def get(self, schema: pyarrow.Schema) -> _ResultSchema: + def get_by_key(self, key: Any) -> Optional[_ResultSchema]: + """The last ``_ResultSchema`` if its kernel schema key equals + ``key``, which costs far less than building and comparing a + ``pyarrow.Schema``.""" + last = self._last + if last is not None and last.key is not None and last.key == key: + return last + return None + + def get(self, schema: pyarrow.Schema, key: Any = None) -> _ResultSchema: last = self._last if last is not None and schema.equals(last.schema, check_metadata=True): + if key is not None: + last.key = key return last - last = self._last = _ResultSchema(schema) + last = self._last = _ResultSchema(schema, key) return last @@ -199,7 +230,9 @@ class KernelResultSet(ResultSet): ``arrow_schema() -> pyarrow.Schema``, ``fetch_next_batch() -> Optional[pyarrow.RecordBatch]``, and ``close() -> None``. Both ``databricks_sql_kernel.ExecutedStatement`` and - ``databricks_sql_kernel.ResultStream`` satisfy that contract. + ``databricks_sql_kernel.ResultStream`` satisfy that contract. Handles + that also have ``schema_key()`` and ``fetch_all_rows()`` let small + results skip pyarrow. """ def __init__( @@ -213,13 +246,19 @@ def __init__( schema_cache: Optional[_SchemaCache] = None, ): try: - schema = kernel_handle.arrow_schema() + key = _schema_key(kernel_handle) + result_schema = None + if key is not None and schema_cache is not None: + result_schema = schema_cache.get_by_key(key) + if result_schema is None: + schema = kernel_handle.arrow_schema() + if schema_cache is None: + result_schema = _ResultSchema(schema, key) + else: + result_schema = schema_cache.get(schema, key) except Exception as exc: raise wrap_kernel_exception("KernelResultSet.arrow_schema", exc) from exc - if schema_cache is None: - self._result_schema = _ResultSchema(schema) - else: - self._result_schema = schema_cache.get(schema) + self._result_schema = result_schema super().__init__( connection=connection, backend=backend, @@ -236,7 +275,18 @@ def __init__( arrow_schema_bytes=None, ) self._kernel_handle = kernel_handle - self._schema: pyarrow.Schema = schema + # Equal to the handle's schema, metadata included. + self._schema: pyarrow.Schema = result_schema.schema + # Kernels with schema keys can also return small results as rows of + # Python values (``fetch_all_rows``), skipping pyarrow. Tried once, + # on the first row fetch unless an Arrow fetch came first, for types + # the pandas-free path converts. + self._try_kernel_rows = ( + key is not None and result_schema.max_rows_without_pandas >= 0 + ) + # Rows of values from ``fetch_all_rows`` and the next one to return. + self._kernel_rows: Optional[List[tuple]] = None + self._kernel_rows_pos = 0 # FIFO of record batches plus a per-head row offset, so # partial fetches (fetchmany(n) for n < batch_size) don't # re-fetch from the kernel. @@ -340,9 +390,69 @@ def _drain_batches(self) -> List[pyarrow.RecordBatch]: self._next_row_index += rows return chunks + # ----- rows converted by the kernel ----- + + def _have_kernel_rows(self) -> bool: + """Whether the rest of the result is in ``_kernel_rows``. The first + call asks the kernel for every remaining row as Python values, which + it returns for results within the pandas-free cost budget; otherwise + the batches it pulled stay in the kernel for the Arrow path.""" + if self._kernel_rows is not None: + return True + if not self._try_kernel_rows or self._buffer or self._exhausted: + return False + self._try_kernel_rows = False + fetch_all_rows = getattr(self._kernel_handle, "fetch_all_rows", None) + if fetch_all_rows is None: + return False + try: + rows = fetch_all_rows( + self._result_schema.max_rows_without_pandas, + self.connection.disable_pandas is True, + ) + except Exception as exc: + raise wrap_kernel_exception("fetch_all_rows", exc) from exc + if rows is None: + return False + self._kernel_rows = rows + self._exhausted = True + self.has_more_rows = False + self.status = CommandState.SUCCEEDED + return True + + def _take_kernel_rows(self, n: Optional[int]) -> List[tuple]: + """Up to ``n`` (all if None) of the remaining kernel rows.""" + rows = cast(List[tuple], self._kernel_rows) + start = self._kernel_rows_pos + end = len(rows) if n is None else min(len(rows), start + n) + self._kernel_rows_pos = end + self._next_row_index += end - start + return rows if start == 0 and end == len(rows) else rows[start:end] + + def _kernel_rows_to_arrow(self, n: Optional[int]) -> pyarrow.Table: + # Only for Arrow fetches after a row fetch: the values convert back to + # the same Arrow values for every type the kernel converts. + rows = self._take_kernel_rows(n) + if not rows: + return pyarrow.Table.from_batches([], schema=self._schema) + return pyarrow.Table.from_arrays( + [ + pyarrow.array(column, type=field.type) + for column, field in zip(zip(*rows), self._schema) + ], + schema=self._schema, + ) + + def _make_rows(self, values: List[tuple]) -> List[Row]: + row_factory = self._result_schema.row_factory + return [row_factory(*v) for v in values] + # ----- Arrow fetches ----- def fetchall_arrow(self) -> pyarrow.Table: + if self._kernel_rows is not None: + return self._kernel_rows_to_arrow(None) + self._try_kernel_rows = False return self._drain() def fetchmany_arrow(self, size: int) -> pyarrow.Table: @@ -350,6 +460,9 @@ def fetchmany_arrow(self, size: int) -> pyarrow.Table: raise ValueError(f"fetchmany_arrow size must be >= 0, got {size}") if size == 0: return pyarrow.Table.from_batches([], schema=self._schema) + if self._kernel_rows is not None: + return self._kernel_rows_to_arrow(size) + self._try_kernel_rows = False self._ensure_buffered(size) return self._take_buffered(size) @@ -365,6 +478,9 @@ def _rows(self, batches: List[pyarrow.RecordBatch]) -> List[Row]: return self._convert_arrow_table(self._to_table(batches)) def fetchone(self) -> Optional[Row]: + if self._have_kernel_rows(): + rows = self._make_rows(self._take_kernel_rows(1)) + return rows[0] if rows else None self._ensure_buffered(1) if self._buffered_rows() == 0: return None @@ -376,10 +492,14 @@ def fetchmany(self, size: int) -> List[Row]: raise ValueError(f"fetchmany size must be >= 0, got {size}") if size == 0: return [] + if self._have_kernel_rows(): + return self._make_rows(self._take_kernel_rows(size)) self._ensure_buffered(size) return self._rows(self._take_buffered_batches(size)) def fetchall(self) -> List[Row]: + if self._have_kernel_rows(): + return self._make_rows(self._take_kernel_rows(None)) return self._rows(self._drain_batches()) def close(self) -> None: @@ -396,6 +516,8 @@ def close(self) -> None: """ if self._kernel_handle is None: return + self._kernel_rows = None + self._try_kernel_rows = False if not self.connection.open: self._kernel_handle = None self._buffer.clear() diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py index beb98efad..5c4f93b27 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/tests/unit/test_kernel_result_set.py @@ -5,7 +5,9 @@ from __future__ import annotations import datetime +import inspect import logging +import math from collections import deque from decimal import Decimal from typing import Deque @@ -22,6 +24,7 @@ from databricks.sql.backend.kernel import result_set as kernel_result_set from databricks.sql.backend.kernel.result_set import KernelResultSet from databricks.sql.backend.types import CommandId, CommandState +from databricks.sql.exc import OperationalError from databricks.sql.result_set import ResultSet @@ -623,11 +626,13 @@ def batch(lo, hi): to_table.assert_called_once() -def test_a_batch_whose_schema_differs_still_fails_as_before(): +@pytest.mark.parametrize("handle_type", ["_FakeKernelHandle", "_FakeRowsHandle"]) +def test_a_batch_whose_schema_differs_still_fails_as_before(handle_type): schema = pa.schema([("n", pa.int64())]) other = pa.RecordBatch.from_arrays([pa.array(["x"])], names=["n"]) + handle = globals()[handle_type](schema, [other]) with pytest.raises(pa.ArrowInvalid): - _make_rs(_FakeKernelHandle(schema, [other])).fetchall() + _make_rs(handle).fetchall() def test_the_cost_budget_is_computed_once_per_schema_not_per_fetch(): @@ -653,3 +658,385 @@ def test_the_cost_budget_is_computed_once_per_schema_not_per_fetch(): _FakeKernelHandle(pa.schema([("m", pa.int64())]), []), schema_cache=cache ) assert budget.call_count == 2 + + +# ----- rows built by the kernel (schema_key / fetch_all_rows) ----- + + +class _FakeSchemaKey: + def __init__(self, schema: pa.Schema): + self.schema = schema + + def __eq__(self, other): + return isinstance(other, _FakeSchemaKey) and self.schema.equals( + other.schema, check_metadata=True + ) + + +class _FakeRowsHandle(_FakeKernelHandle): + """A handle from a kernel that has ``schema_key`` and ``fetch_all_rows``, + which returns the rows as ``to_pylist()`` values when every remaining + row fits in ``max_rows``, and otherwise keeps the batches it pulled for + ``fetch_next_batch``, as the kernel does.""" + + def __init__(self, schema, batches, convertible=True): + super().__init__(schema, batches) + self._pending: Deque[pa.RecordBatch] = deque() + self._convertible = convertible + self.schema_calls = 0 + self.rows_calls = [] + + def arrow_schema(self): + self.schema_calls += 1 + return super().arrow_schema() + + def schema_key(self): + return _FakeSchemaKey(self._schema) + + def fetch_all_rows(self, max_rows, nan_ok=False): + self.rows_calls.append((max_rows, nan_ok)) + while sum(b.num_rows for b in self._pending) <= max_rows and self._batches: + self._pending.append(self._batches.popleft()) + if self._batches or not self._convertible: + return None + if sum(b.num_rows for b in self._pending) > max_rows: + return None + # The kernel declines batches whose column types differ. + if any(b.schema.types != self._schema.types for b in self._pending): + return None + columns = [[] for _ in self._schema] + for batch in self._pending: + for values, column in zip(columns, batch.columns): + values.extend(column.to_pylist()) + if not nan_ok and any( + isinstance(v, float) and v != v for values in columns for v in values + ): + return None + self._pending.clear() + return list(zip(*columns)) + + def fetch_next_batch(self): + if self._pending: + return self._pending.popleft() + return super().fetch_next_batch() + + +_ROWS_SCHEMA = pa.schema( + [ + pa.field("n", pa.int64(), metadata={b"databricks.type_name": b"BIGINT"}), + pa.field("d", pa.decimal128(10, 2)), + pa.field("t", pa.timestamp("us")), + pa.field("s", pa.string()), + ] +) + + +def _rows_batch(lo, hi, schema=_ROWS_SCHEMA): + return pa.RecordBatch.from_arrays( + [ + pa.array(range(lo, hi), schema.field(0).type), + pa.array( + [Decimal(i).scaleb(-2) for i in range(lo, hi)], schema.field(1).type + ), + pa.array([i * 3_600_000_000 for i in range(lo, hi)], schema.field(2).type), + pa.array([None if i % 2 else str(i) for i in range(lo, hi)]), + ], + schema=schema, + ) + + +def _expected_rows(*batches): + table = pa.Table.from_batches(batches) + rs = _make_rs(_FakeKernelHandle(table.schema, [])) + return ResultSet._convert_arrow_table(rs, table) + + +def test_kernel_rows_skip_pyarrow_for_schema_and_values(): + cache = kernel_result_set._SchemaCache() + first = _FakeRowsHandle(_reimported(_ROWS_SCHEMA), [_rows_batch(0, 3)]) + second = _FakeRowsHandle(_reimported(_ROWS_SCHEMA), [_rows_batch(3, 5)]) + with patch.object(_FakeRowsHandle, "fetch_next_batch") as next_batch: + rows = _make_rs(first, schema_cache=cache).fetchall() + rs = _make_rs(second, schema_cache=cache) + assert ( + rs.description == _make_rs(_FakeKernelHandle(_ROWS_SCHEMA, [])).description + ) + rows += rs.fetchall() + next_batch.assert_not_called() + # The second statement's schema came from the cache, by key. + assert (first.schema_calls, second.schema_calls) == (1, 0) + _assert_same_rows(rows, _expected_rows(_rows_batch(0, 3), _rows_batch(3, 5))) + assert rows[0].n == 0 and rs.rownumber == 2 + assert rs.status == CommandState.SUCCEEDED and not rs.has_more_rows + + +def test_kernel_rows_with_a_time_zone_stay_within_the_budget(): + schema = pa.schema([("t", pa.timestamp("us", tz="America/Los_Angeles"))]) + budget = kernel_result_set._max_rows_without_pandas(schema.types) + + def batch(n): + return pa.RecordBatch.from_arrays( + [ + pa.array( + range(0, n * 3_600_000_000, 3_600_000_000), schema.field(0).type + ) + ], + schema=schema, + ) + + handle = _FakeRowsHandle(schema, [batch(budget)]) + _assert_same_rows(_make_rs(handle).fetchall(), _expected_rows(batch(budget))) + assert handle.rows_calls == [(budget, False)] + + +def test_a_different_schema_key_is_derived_again(): + cache = kernel_result_set._SchemaCache() + _make_rs(_FakeRowsHandle(_ROWS_SCHEMA, []), schema_cache=cache) + other = pa.schema([f.remove_metadata() for f in _ROWS_SCHEMA]) + handle = _FakeRowsHandle(other, []) + rs = _make_rs(handle, schema_cache=cache) + assert handle.schema_calls == 1 + assert rs.description == _make_rs(_FakeKernelHandle(other, [])).description + + +def test_row_fetches_are_served_from_one_kernel_call(): + handle = _FakeRowsHandle(_ROWS_SCHEMA, [_rows_batch(0, 2), _rows_batch(2, 5)]) + rs = _make_rs(handle) + rows = [rs.fetchone()] + rs.fetchmany(2) + rs.fetchmany(0) + rs.fetchall() + assert rs.fetchone() is None and rs.fetchmany(3) == [] and rs.fetchall() == [] + budget = kernel_result_set._max_rows_without_pandas(_ROWS_SCHEMA.types) + assert handle.rows_calls == [(budget, False)] + _assert_same_rows(rows, _expected_rows(_rows_batch(0, 2), _rows_batch(2, 5))) + assert rs.rownumber == 5 + assert [tuple(r) for r in rs] == [] + + +@pytest.mark.parametrize("disable_pandas", [False, True]) +def test_kernel_keeps_nan_only_without_pandas(disable_pandas): + schema = pa.schema([("f", pa.float64())]) + batch = pa.RecordBatch.from_arrays([pa.array([float("nan"), 1.0])], schema=schema) + handle = _FakeRowsHandle(schema, [batch]) + rs = _make_rs(handle, disable_pandas=disable_pandas) + rows = rs.fetchall() + assert handle.rows_calls[0][1] is disable_pandas + if disable_pandas: + # as_py(), as without kernel rows. + assert isinstance(rows[0][0], float) and math.isnan(rows[0][0]) + assert rows[1][0] == 1.0 + else: + # The kernel declines NaN, so pandas decides it, as before. + _assert_same_rows(rows, ResultSet._convert_arrow_table(rs, pa.table(batch))) + + +def test_results_the_kernel_declines_come_from_the_batches_it_pulled(): + budget = kernel_result_set._max_rows_without_pandas(_ROWS_SCHEMA.types) + batches = [ + _rows_batch(0, 2), + _rows_batch(2, budget), + _rows_batch(budget, budget + 1), + ] + expected = _expected_rows(*batches) + handle = _FakeRowsHandle(_ROWS_SCHEMA, batches) + rs = _make_rs(handle) + rows = [rs.fetchone()] + # Declined after pulling every batch; they are now in the kernel's queue. + assert len(handle.rows_calls) == 1 and not handle._batches + rows += rs.fetchmany(3) + rs.fetchall() + assert len(handle.rows_calls) == 1 + _assert_same_rows(rows, expected) + assert rs.rownumber == budget + 1 + with _spy_to_table() as to_table: + handle = _FakeRowsHandle(_ROWS_SCHEMA, batches) + _assert_same_rows(_make_rs(handle).fetchall(), expected) + to_table.assert_called_once() + + +def test_schemas_the_kernel_cannot_convert_never_ask_it_for_rows(): + schema = pa.schema([("l", pa.list_(pa.int64()))]) + batch = pa.RecordBatch.from_arrays([pa.array([[1], None])], schema=schema) + handle = _FakeRowsHandle(schema, [batch]) + assert [r[0] for r in _make_rs(handle).fetchall()] == [[1], None] + assert handle.rows_calls == [] + + +def test_unconverted_rows_fall_back_to_the_batches(): + handle = _FakeRowsHandle(_ROWS_SCHEMA, [_rows_batch(0, 3)], convertible=False) + rs = _make_rs(handle) + _assert_same_rows(rs.fetchall(), _expected_rows(_rows_batch(0, 3))) + assert len(handle.rows_calls) == 1 + + +@pytest.mark.parametrize( + "fetch", + [ + lambda rs: rs.fetchall_arrow(), + lambda rs: pa.concat_tables( + [rs.fetchmany_arrow(1), rs.fetchmany_arrow(0), rs.fetchmany_arrow(10)] + ), + ], + ids=["fetchall_arrow", "fetchmany_arrow"], +) +def test_arrow_fetches_after_a_row_fetch_return_the_remaining_rows(fetch): + rs = _make_rs(_FakeRowsHandle(_ROWS_SCHEMA, [_rows_batch(0, 4)])) + assert rs.fetchone()[0] == 0 + table = fetch(rs) + assert table.schema.equals(_ROWS_SCHEMA, check_metadata=True) + assert table.equals(pa.Table.from_batches([_rows_batch(1, 4)])) + assert rs.fetchall_arrow().num_rows == 0 and rs.rownumber == 4 + + +@pytest.mark.parametrize( + "name, disable_pandas", + [(name, False) for name in sorted(_columns_skipping_pandas())] + + [("f64_nan", True), ("f32_nan", True)], +) +def test_every_kernel_converted_type_round_trips_to_arrow(name, disable_pandas): + columns = {**_columns_skipping_pandas(), **_columns_needing_pandas()} + batch = pa.record_batch([pa.array(range(4)), columns[name]], names=["k", name]) + with patch.object(kernel_result_set, "_MAX_COST_WITHOUT_PANDAS", 10**9): + handle = _FakeRowsHandle(batch.schema, [batch]) + rs = _make_rs(handle, disable_pandas=disable_pandas) + assert rs.fetchone()[0] == 0 + table = rs.fetchall_arrow() + assert handle.rows_calls + expected = pa.Table.from_batches([batch.slice(1)]) + assert table.schema.equals(expected.schema, check_metadata=True) + # repr: NaN != NaN for Table.equals. + assert repr(table.to_pylist()) == repr(expected.to_pylist()) + assert table.equals(expected) or disable_pandas + + +class _KeyOnlyHandle(_FakeKernelHandle): + def schema_key(self): + return _FakeSchemaKey(self._schema) + + +def test_a_failing_schema_key_falls_back_to_the_arrow_schema(caplog): + handle = _FakeRowsHandle(_ROWS_SCHEMA, [_rows_batch(0, 3)]) + handle.schema_key = MagicMock(side_effect=RuntimeError("boom")) + with caplog.at_level(logging.DEBUG, logger=kernel_result_set.__name__): + rs = _make_rs(handle, schema_cache=kernel_result_set._SchemaCache()) + assert "schema_key() failed" in caplog.text + assert rs.description == _make_rs(_FakeKernelHandle(_ROWS_SCHEMA, [])).description + _assert_same_rows(rs.fetchall(), _expected_rows(_rows_batch(0, 3))) + assert handle.rows_calls == [] + + +def test_handles_with_schema_keys_but_without_fetch_all_rows_use_the_batches(): + rs = _make_rs(_KeyOnlyHandle(_ROWS_SCHEMA, [_rows_batch(0, 3)])) + _assert_same_rows(rs.fetchall(), _expected_rows(_rows_batch(0, 3))) + + +@pytest.mark.parametrize("taken", [1, 3], ids=["mid_batch", "batch_boundary"]) +def test_row_fetches_after_an_arrow_fetch_stay_on_the_arrow_path(taken): + handle = _FakeRowsHandle(_ROWS_SCHEMA, [_rows_batch(0, 3), _rows_batch(3, 4)]) + rs = _make_rs(handle) + assert rs.fetchmany_arrow(taken).num_rows == taken + rows = rs.fetchall() + assert handle.rows_calls == [] + _assert_same_rows(rows, _expected_rows(_rows_batch(taken, 4))) + + +def test_kernel_row_errors_are_mapped(): + handle = _FakeRowsHandle(_ROWS_SCHEMA, [_rows_batch(0, 1)]) + handle.fetch_all_rows = MagicMock(side_effect=RuntimeError("boom")) + with pytest.raises(OperationalError, match="fetch_all_rows"): + _make_rs(handle).fetchall() + + +def test_a_fetch_after_a_kernel_row_error_returns_every_row_once(): + batches = [_rows_batch(0, 2), _rows_batch(2, 4), _rows_batch(4, 5)] + handle = _FakeRowsHandle(_ROWS_SCHEMA, batches) + + def pull_two_then_fail(max_rows, nan_ok=False): + # As the kernel does when reading fails mid-way: the batches read so + # far stay queued for fetch_next_batch. + handle._pending.extend([handle._batches.popleft() for _ in range(2)]) + raise RuntimeError("connection reset") + + handle.fetch_all_rows = pull_two_then_fail + rs = _make_rs(handle) + with pytest.raises(OperationalError): + rs.fetchone() + rows = [rs.fetchone()] + rs.fetchall() + _assert_same_rows(rows, _expected_rows(*batches)) + assert rs.rownumber == 5 + + +def test_close_drops_kernel_rows(): + rs = _make_rs(_FakeRowsHandle(_ROWS_SCHEMA, [_rows_batch(0, 3)])) + assert rs.fetchone()[0] == 0 + rs.close() + assert rs.fetchall() == [] + + +@pytest.mark.realkernel +class TestKernelRowsMatchThePandasPath: + """The real kernel's conversion (``_rows_from_arrow`` runs the one behind + ``fetch_all_rows``) returns the values of the pandas path for every type + that skips pandas. Skipped for kernels without it.""" + + @staticmethod + def _rows_from_arrow(): + kernel = pytest.importorskip("databricks_sql_kernel") + if not hasattr(kernel.ExecutedStatement, "fetch_all_rows"): + pytest.skip("kernel has no fetch_all_rows") + rows_from_arrow = getattr(kernel, "_rows_from_arrow", None) + assert rows_from_arrow is not None, "kernel lost its conversion test hook" + return rows_from_arrow + + @pytest.mark.parametrize("name", sorted(_columns_skipping_pandas())) + def test_values(self, name): + rows_from_arrow = self._rows_from_arrow() + column = _columns_skipping_pandas()[name] + table = pa.table({"k": pa.array(range(len(column))), name: column}) + rs = _make_rs(_FakeKernelHandle(table.schema, [])) + expected = ResultSet._convert_arrow_table(rs, table) + values = rows_from_arrow(table.to_batches()) + _assert_same_rows([rs._result_schema.row_factory(*v) for v in values], expected) + + @pytest.mark.parametrize("name", sorted(_columns_needing_pandas())) + def test_types_needing_pandas_are_declined(self, name): + rows_from_arrow = self._rows_from_arrow() + table = pa.table({name: _columns_needing_pandas()[name]}) + assert rows_from_arrow(table.to_batches()) is None + + def test_handles_take_the_arguments_this_module_passes(self): + self._rows_from_arrow() + import databricks_sql_kernel as kernel + + for handle_type in (kernel.ExecutedStatement, kernel.ResultStream): + params = list( + inspect.signature(handle_type.fetch_all_rows).parameters.values() + ) + assert [(p.name, p.kind) for p in params[1:]] == [ + ("max_rows", inspect.Parameter.POSITIONAL_OR_KEYWORD), + ("nan_ok", inspect.Parameter.POSITIONAL_OR_KEYWORD), + ] + assert callable(handle_type.schema_key) + schema = pa.schema( + [pa.field("n", pa.int64(), metadata={b"databricks.type_name": b"BIGINT"})] + ) + key = kernel.SchemaKey(schema) + assert key == kernel.SchemaKey(_reimported(schema)) + assert key != kernel.SchemaKey(pa.schema([pa.field("n", pa.int64())])) + # Keys tell apart everything Schema.equals(check_metadata=True) does. + for other in ( + schema.field(0).with_nullable(False), + schema.field(0).with_name("m"), + schema.field(0).with_type(pa.int32()), + ): + assert not schema.equals(pa.schema([other]), check_metadata=True) + assert key != kernel.SchemaKey(pa.schema([other])) + # The conversion declines a batch whose column types differ. + batches = [ + pa.record_batch([pa.array([1])], names=["n"]), + pa.record_batch([pa.array(["x"])], names=["n"]), + ] + assert self._rows_from_arrow()(batches) is None + cache = kernel_result_set._SchemaCache() + cache.get(schema, key) + assert cache.get_by_key(kernel.SchemaKey(_reimported(schema))) is not None + assert ( + cache.get_by_key(kernel.SchemaKey(pa.schema([("m", pa.int64())]))) is None + )