diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bae4c54e..e62784f4c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ - Fix the kernel backend ignoring `use_cloud_fetch=False`, so complete results can be fetched without cloud-storage downloads. - Fix inline parameters (`use_inline_params=True`) so a sequence directly after `IN` renders as a list again: `x IN %(ids)s` renders as `x IN (1,2,3)` instead of `x IN ARRAY(1,2,3)`, which Spark rejects. Inside the list, tuples render as rows for a multi-column `IN`, and a `set` or `frozenset` is accepted again. Sequences anywhere else still render as `ARRAY(...)` ([#629](https://github.com/databricks/databricks-sql-python/issues/629)) - 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. # 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 15f78bf5f..fcb798158 100644 --- a/src/databricks/sql/backend/kernel/result_set.py +++ b/src/databricks/sql/backend/kernel/result_set.py @@ -32,7 +32,7 @@ import logging from collections import deque -from typing import Any, Deque, List, Optional, TYPE_CHECKING, cast +from typing import Any, Deque, List, Optional, Tuple, TYPE_CHECKING, cast import pyarrow @@ -48,6 +48,109 @@ logger = logging.getLogger(__name__) +# Small results are converted with to_pylist() instead of through a pandas +# DataFrame, whose fixed cost (100-300 µs) dominates them. to_pylist() costs +# more per value, so a result skips pandas only while its row count times the +# sum of its column weights (per-value cost relative to an integer) is within +# _MAX_COST_WITHOUT_PANDAS, which kept it at least 25% cheaper than pandas for +# every type and width measured (pyarrow 23, pandas 1.3-3.1). +_MAX_ROWS_WITHOUT_PANDAS = 128 +_MAX_COST_WITHOUT_PANDAS = 1024 + +# Types for which to_pylist() returns the same values as the pandas path in +# ResultSet._convert_arrow_table, with their weights. Floats only qualify +# without NaN, which pandas turns into None or keeps depending on its version; +# timestamps only when not in nanoseconds. +_FLOAT_TYPE_IDS = frozenset(t.id for t in (pyarrow.float32(), pyarrow.float64())) +_TIMESTAMP_TYPE_ID = pyarrow.timestamp("us").id +_VALUE_WEIGHTS = { + **{ + t.id: 1 + for t in ( + pyarrow.int8(), + pyarrow.int16(), + pyarrow.int32(), + pyarrow.int64(), + pyarrow.uint8(), + pyarrow.uint16(), + pyarrow.uint32(), + pyarrow.uint64(), + pyarrow.bool_(), + pyarrow.float32(), + pyarrow.float64(), + pyarrow.null(), + ) + }, + **{t.id: 3 for t in (pyarrow.string(), pyarrow.large_string())}, + **{ + t.id: 6 + for t in ( + pyarrow.binary(), + pyarrow.large_binary(), + pyarrow.binary(1), + pyarrow.decimal128(1), + pyarrow.decimal256(1), + ) + }, + pyarrow.date32().id: 12, + _TIMESTAMP_TYPE_ID: 12, +} +# pyarrow looks the time zone up again for every tz-aware value, which costs +# tens of microseconds when pytz isn't installed. +_TZ_TIMESTAMP_WEIGHT = 400 + + +def _max_rows_without_pandas(types: List[pyarrow.DataType]) -> int: + """The most rows of a result with these column types that are converted + with to_pylist(); -1 if such results always go through pandas.""" + total = 0 + for arrow_type in types: + weight = _VALUE_WEIGHTS.get(arrow_type.id) + if weight is None: + return -1 + if arrow_type.id == _TIMESTAMP_TYPE_ID: + if arrow_type.unit == "ns": + return -1 + if arrow_type.tz is not None: + weight = _TZ_TIMESTAMP_WEIGHT + total += weight + if total == 0: + return -1 + return min(_MAX_ROWS_WITHOUT_PANDAS, _MAX_COST_WITHOUT_PANDAS // total) + + +def _rows_without_pandas( + table: pyarrow.Table, description: List[Tuple] +) -> Optional[List[Row]]: + """The rows of ``table`` built from ``to_pylist()``, or None if the table + must go through pandas (too costly, or a column whose values could + differ).""" + num_rows = table.num_rows + # No weight is below 1, so this rules out larger results without + # looking at every column type. + if num_rows * table.num_columns > _MAX_COST_WITHOUT_PANDAS: + return None + if num_rows > _max_rows_without_pandas(table.schema.types): + return None + columns = [] + for column in table.itercolumns(): + try: + values = column.to_pylist() + except Exception: + # E.g. a date past 9999-12-31: pandas then raises its own error, + # as it does without this fast path. + logger.debug( + "to_pylist() failed for a %s column; using pandas", + column.type, + exc_info=True, + ) + return None + if column.type.id in _FLOAT_TYPE_IDS and any(v != v for v in values): + return None + columns.append(values) + row_factory = Row(*[c[0] for c in description]) + return [row_factory(*values) for values in zip(*columns)] + class KernelResultSet(ResultSet): """Streaming ``ResultSet`` over a kernel handle. @@ -202,6 +305,13 @@ def fetchmany_arrow(self, size: int) -> pyarrow.Table: # ----- Row fetches ----- + def _convert_arrow_table(self, table): + if self.connection.disable_pandas is not True: + rows = _rows_without_pandas(table, self.description) + if rows is not None: + return rows + return super()._convert_arrow_table(table) + def fetchone(self) -> Optional[Row]: self._ensure_buffered(1) if self._buffered_rows() == 0: diff --git a/tests/unit/test_kernel_result_set.py b/tests/unit/test_kernel_result_set.py index cfa2d3f2e..aafb45679 100644 --- a/tests/unit/test_kernel_result_set.py +++ b/tests/unit/test_kernel_result_set.py @@ -4,10 +4,14 @@ from __future__ import annotations +import datetime +import logging from collections import deque +from decimal import Decimal from typing import Deque -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch +import numpy import pytest # pyarrow is an optional connector dep; the default-deps CI test @@ -15,8 +19,10 @@ # so the whole module must skip when pyarrow is unavailable. pa = pytest.importorskip("pyarrow") +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.result_set import ResultSet class _FakeKernelHandle: @@ -243,3 +249,271 @@ def test_close_skips_kernel_call_when_connection_already_closed(int_schema): # …but local state is still terminal so __del__ is idempotent: assert rs.status == CommandState.CLOSED assert rs.has_been_closed_server_side is True + + +# ----- small results converted without pandas ----- + +_UTC = datetime.timezone.utc + + +def _columns_skipping_pandas(): + return { + "i8": pa.array([1, None, -128, 127], pa.int8()), + "i16": pa.array([1, None, -(2**15), 2**15 - 1], pa.int16()), + "i32": pa.array([1, None, -(2**31), 2**31 - 1], pa.int32()), + "i64": pa.array([1, None, -(2**63), 2**63 - 1], pa.int64()), + "u8": pa.array([1, None, 0, 255], pa.uint8()), + "u16": pa.array([1, None, 0, 2**16 - 1], pa.uint16()), + "u32": pa.array([1, None, 0, 2**32 - 1], pa.uint32()), + "u64": pa.array([1, None, 0, 2**64 - 1], pa.uint64()), + "bool": pa.array([True, None, False, True]), + "f32": pa.array([0.1, None, float("inf"), -0.0], pa.float32()), + "f64": pa.array([0.1, None, float("-inf"), -0.0], pa.float64()), + "str": pa.array(["a", None, "", "é\U0001f600"]), + "large_str": pa.array(["a", None, "", "z"], pa.large_string()), + "bin": pa.array([b"a", None, b"", b"\x00\xff"]), + "large_bin": pa.array([b"a", None, b"", b"\x00"], pa.large_binary()), + "fixed_bin": pa.array([b"ab", None, b"cd", b"\x00\x00"], pa.binary(2)), + "dec": pa.array( + [Decimal("123141078.2283"), None, Decimal("-0.0001"), Decimal("0.0000")], + pa.decimal128(38, 4), + ), + "dec256": pa.array( + [Decimal("1.5"), None, Decimal("-2.5"), Decimal("0.0")], + pa.decimal256(40, 1), + ), + "date": pa.array( + [ + datetime.date(2024, 1, 1), + None, + datetime.date(1, 1, 1), + datetime.date(9999, 12, 31), + ] + ), + "ts": pa.array( + [ + datetime.datetime(2024, 1, 1, 1, 2, 3, 4), + None, + datetime.datetime(1, 1, 1), + datetime.datetime(9999, 12, 31, 23, 59, 59, 999999), + ], + pa.timestamp("us"), + ), + "ts_s": pa.array([0, None, -1, 2], pa.timestamp("s")), + "ts_ms": pa.array([0, None, 1, 2], pa.timestamp("ms")), + "ts_utc": pa.array( + [ + datetime.datetime(2024, 1, 1, tzinfo=_UTC), + None, + datetime.datetime(1, 1, 1, tzinfo=_UTC), + datetime.datetime(2024, 3, 10, 10, tzinfo=_UTC), + ], + pa.timestamp("us", tz="Etc/UTC"), + ), + # Includes the instant when DST starts. + "ts_la": pa.array( + [0, None, 1710064800000000, 1699174800000000], + pa.timestamp("us", tz="America/Los_Angeles"), + ), + "ts_offset": pa.array([0, None, 1, 2], pa.timestamp("us", tz="+05:30")), + "null": pa.nulls(4), + } + + +def test_every_type_skipping_pandas_has_a_column_above(): + covered = {column.type.id for column in _columns_skipping_pandas().values()} + assert covered == set(kernel_result_set._VALUE_WEIGHTS) + + +def _columns_needing_pandas(): + return { + # pandas maps NaN to None or keeps it, depending on its version. + "f64_nan": pa.array([float("nan"), None, 1.0, 2.0]), + "f32_nan": pa.array([1.0, float("nan"), None, 2.0], pa.float32()), + "ts_ns": pa.array([1, None, 1000, 2], pa.timestamp("ns")), + # Out-of-range values convert differently; see + # test_out_of_range_values_convert_as_on_the_pandas_path. + "date64": pa.array([0, None, 86_400_000, -1], pa.date64()), + "time32_s": pa.array([0, None, 1, 2], pa.time32("s")), + "time32_ms": pa.array([0, None, 1, 2], pa.time32("ms")), + "time64": pa.array([0, None, 1, 86_399_999_999], pa.time64("us")), + "time_ns": pa.array([1000, None, 2000, 0], pa.time64("ns")), + "duration": pa.array([0, None, 1, 2], pa.duration("us")), + "interval": pa.array( + [pa.MonthDayNano([1, 2, 3]), None, pa.MonthDayNano([0, 0, 0]), None], + pa.month_day_nano_interval(), + ), + "list": pa.array([[1, 2], None, [], [None]], pa.list_(pa.int64())), + "struct": pa.array( + [{"x": 1}, None, {"x": None}, {"x": 2}], pa.struct([("x", pa.int64())]) + ), + "map": pa.array( + [[("a", 1)], None, [], [("b", None)]], pa.map_(pa.string(), pa.int64()) + ), + "dictionary": pa.array(["a", None, "a", "b"]).dictionary_encode(), + "f16": pa.array( + numpy.array([1, 0, 2, 3], numpy.float16), + mask=numpy.array([False, True, False, False]), + ), + } + + +def _assert_same_rows(rows, expected): + assert len(rows) == len(expected) + for row, expected_row in zip(rows, expected): + assert row.__fields__ == expected_row.__fields__ + assert len(row) == len(expected_row) + for value, expected_value in zip(row, expected_row): + assert type(value) is type(expected_value) + assert repr(value) == repr(expected_value) + if isinstance(value, datetime.datetime): + assert type(value.tzinfo) is type(expected_value.tzinfo) + assert value.tzinfo == expected_value.tzinfo + + +def _convert(table): + """(rows from KernelResultSet, rows from the pandas path, whether the + kernel result set used the pandas path).""" + rs = _make_rs(_FakeKernelHandle(table.schema, [])) + pandas_path = ResultSet._convert_arrow_table + with patch.object( + ResultSet, "_convert_arrow_table", autospec=True, side_effect=pandas_path + ) as spy: + rows = rs._convert_arrow_table(table) + return rows, pandas_path(rs, table), spy.called + + +@pytest.mark.parametrize("name", sorted(_columns_skipping_pandas())) +def test_small_results_skip_pandas_with_the_same_values(name): + column = _columns_skipping_pandas()[name] + tables = [ + pa.table({name: column}), + pa.table({name: column.slice(0, 1)}), + pa.table({name: column.slice(1, 1)}), + # Chunked, next to other columns, with a duplicate column name. + pa.table( + [ + pa.chunked_array([[1, 2], [3, 4]]), + pa.chunked_array([column.slice(0, 1), column.slice(1)]), + pa.array(["a", "b", None, "d"]), + ], + names=["k", name, "k"], + ), + ] + # Values only; the cost budget is tested separately. + with patch.object(kernel_result_set, "_MAX_COST_WITHOUT_PANDAS", 10**9): + for table in tables: + rows, expected, used_pandas = _convert(table) + assert not used_pandas + _assert_same_rows(rows, expected) + + +@pytest.mark.parametrize("name", sorted(_columns_needing_pandas())) +def test_results_whose_values_could_differ_use_pandas(name): + table = pa.table( + {"k": pa.array([1, 2, 3, 4]), name: _columns_needing_pandas()[name]} + ) + rows, expected, used_pandas = _convert(table) + assert used_pandas + assert len(rows) == len(expected) == 4 + + +@pytest.mark.parametrize( + "types", + [ + [pa.int64()], + [pa.float64()] * 16, + [pa.string()] * 3, + [pa.decimal128(38, 4), pa.int32(), pa.string(), pa.date32()], + [pa.timestamp("us", "Etc/UTC")], + [pa.timestamp("us"), pa.binary()] * 50, + ], + ids=["int", "16_doubles", "3_strings", "mixed", "tz_timestamp", "wide"], +) +def test_results_over_the_cost_budget_use_pandas(types): + limit = kernel_result_set._max_rows_without_pandas(types) + assert 0 < limit <= kernel_result_set._MAX_ROWS_WITHOUT_PANDAS + for num_rows, expect_pandas in ((limit, False), (limit + 1, True)): + table = pa.table( + [pa.nulls(num_rows, t) for t in types], + names=[f"c{i}" for i in range(len(types))], + ) + rows, expected, used_pandas = _convert(table) + assert used_pandas is expect_pandas + _assert_same_rows(rows, expected) + + +def test_cost_budget_scales_with_width_and_type(): + max_rows = kernel_result_set._max_rows_without_pandas + assert max_rows([pa.int64()]) == kernel_result_set._MAX_ROWS_WITHOUT_PANDAS + assert max_rows([pa.int64()] * 64) < max_rows([pa.int64()] * 8) + assert max_rows([pa.date32()] * 8) < max_rows([pa.int64()] * 8) + assert max_rows([pa.timestamp("us", "UTC")]) < max_rows([pa.timestamp("us")]) + # Too costly for even one row; only empty results skip pandas. + assert max_rows([pa.timestamp("us", "UTC")] * 3) == 0 + assert max_rows([]) == -1 + assert max_rows([pa.int64(), pa.list_(pa.int64())]) == -1 + assert max_rows([pa.timestamp("ns")]) == -1 + + +@pytest.mark.parametrize( + "column, expected", + [ + # Before 1970 and not whole days: to_pylist() rounds down. + (pa.array([-1], pa.date64()), [datetime.date(1970, 1, 1)]), + # Outside 0-24h: to_pylist() wraps around. + (pa.array([86_400], pa.time32("s")), ValueError), + (pa.array([-1], pa.time64("us")), ValueError), + # Past datetime's range: to_pylist() raises OverflowError instead. + (pa.array([2_932_897], pa.date32()), ValueError), + (pa.array([253_402_300_800_000_000], pa.timestamp("us", "UTC")), ValueError), + ], + ids=["date64", "time32", "time64", "date32", "timestamp"], +) +def test_out_of_range_values_convert_as_on_the_pandas_path(column, expected): + table = pa.table({"c": column}) + rs = _make_rs(_FakeKernelHandle(table.schema, [])) + for convert in ( + rs._convert_arrow_table, + lambda t: ResultSet._convert_arrow_table(rs, t), + ): + if isinstance(expected, list): + assert [row[0] for row in convert(table)] == expected + else: + with pytest.raises(expected): + convert(table) + + +def test_conversion_errors_are_logged_before_falling_back(caplog): + table = pa.table({"c": pa.array([2_932_897], pa.date32())}) + rs = _make_rs(_FakeKernelHandle(table.schema, [])) + with caplog.at_level(logging.DEBUG, logger=kernel_result_set.__name__): + with pytest.raises(ValueError): + rs._convert_arrow_table(table) + assert "to_pylist() failed for a date32[day] column" in caplog.text + + +def test_empty_results_skip_pandas(): + rows, expected, used_pandas = _convert(pa.table({"n": pa.array([], pa.int64())})) + assert rows == expected == [] + assert not used_pandas + + +def test_fetch_methods_return_the_pandas_path_values(): + schema = pa.schema([("d", pa.decimal128(38, 4)), ("t", pa.timestamp("us"))]) + assert kernel_result_set._max_rows_without_pandas(schema.types) >= 3 + batch = pa.RecordBatch.from_arrays( + [ + pa.array( + [Decimal("1.5000"), None, Decimal("-2.0001")], schema.field(0).type + ), + pa.array([0, 1, None], schema.field(1).type), + ], + schema=schema, + ) + expected = ResultSet._convert_arrow_table( + _make_rs(_FakeKernelHandle(schema, [])), pa.Table.from_batches([batch]) + ) + _assert_same_rows(_make_rs(_FakeKernelHandle(schema, [batch])).fetchall(), expected) + rs = _make_rs(_FakeKernelHandle(schema, [batch])) + _assert_same_rows([rs.fetchone()] + rs.fetchmany(5), expected)