diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bae4c54e..bc133f765 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,8 @@ - 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 the first `connect()` in a process by importing pandas only when rows are first converted through it. +- Speed up the first `connect()` with `use_kernel=True` by importing pyarrow (and numpy, whose BLAS threads compete for CPU while connections open) on the first result instead of in `connect()`. # 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/client.py b/src/databricks/sql/backend/kernel/client.py index df52f34a5..d3e697eb7 100644 --- a/src/databricks/sql/backend/kernel/client.py +++ b/src/databricks/sql/backend/kernel/client.py @@ -45,8 +45,6 @@ wrap_kernel_exception as _wrap_kernel_exception, ) from databricks.sql.backend.kernel.auth_bridge import kernel_auth_kwargs -from databricks.sql.backend.kernel.result_set import KernelResultSet -from databricks.sql.backend.kernel.type_mapping import bind_tspark_params from databricks.sql.backend.types import ( BackendType, CommandId, @@ -664,6 +662,10 @@ def execute_command( # connector, which the kernel accepts directly. stmt.set_query_tags(query_tags) if parameters: + from databricks.sql.backend.kernel.type_mapping import ( + bind_tspark_params, + ) + bind_tspark_params(stmt, parameters) if async_op: async_exec = stmt.submit() @@ -1048,6 +1050,9 @@ def _make_result_set( ``cursor.query_id`` / ``get_query_state`` would stay pinned to a prior query after a metadata call (the metadata methods mint a synthetic command id but previously never published it).""" + # Imported on first result, not at connect: it imports pyarrow. + from databricks.sql.backend.kernel.result_set import KernelResultSet + cursor.active_command_id = command_id return KernelResultSet( connection=cursor.connection, diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 613b62386..cb7fc0b18 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -12,12 +12,13 @@ BinaryIO, TYPE_CHECKING, ) -import pandas -try: +from databricks.sql.common.optional_pyarrow import optional_pyarrow + +if TYPE_CHECKING: import pyarrow -except ImportError: - pyarrow = None +else: + pyarrow = optional_pyarrow(globals()) import json import os import decimal diff --git a/src/databricks/sql/common/optional_pyarrow.py b/src/databricks/sql/common/optional_pyarrow.py new file mode 100644 index 000000000..b67d28b5a --- /dev/null +++ b/src/databricks/sql/common/optional_pyarrow.py @@ -0,0 +1,75 @@ +"""``pyarrow`` for modules that treat it as optional, imported on first use. + +``connect()`` imports these modules, but pyarrow is only needed once results are +read. Importing it also imports numpy, which starts OpenBLAS's thread pool (one +busy-waiting thread per core), so it shouldn't be part of connecting. + +One case differs from importing pyarrow eagerly: when pyarrow is installed but +fails to import (e.g. a numpy ABI mismatch). These modules then treat it as +available and fail when an attribute is first used, while the Thrift and SEA +backends, which import pyarrow when they load, treat it as missing as before. +""" + +import importlib +import importlib.util +import sys +import types +from typing import Any, Dict, List, Optional + +# Module globals whose ``pyarrow`` is the stand-in, rebound on first use. +_NAMESPACES: List[Dict[str, Any]] = [] + + +class _LazyPyarrow(types.ModuleType): + """Stands in for the ``pyarrow`` module until an attribute is used. + + The first use rebinds ``pyarrow`` to the real module in every registered + namespace still holding the stand-in, so reads on the hot path afterwards + are plain module attribute lookups. The stand-in keeps no copies: a read + through it always gets the live module's attribute, patches included. + """ + + def __getattr__(self, name: str) -> Any: + # A failed import is tried once, as an eager import was: retrying + # would rerun pyarrow's (and numpy's) partial initialization each time. + error = self.__dict__.get("_import_error") + if error is None: + try: + module = importlib.import_module("pyarrow") + except ImportError as e: + self.__dict__["_import_error"] = error = e + if error is not None: + # AttributeError, so hasattr() and getattr(..., default) still work. + raise AttributeError(f"pyarrow could not be imported: {error}") from error + for namespace in _NAMESPACES: + if namespace.get("pyarrow") is self: + namespace["pyarrow"] = module + return getattr(module, name) + + +_LAZY_PYARROW = _LazyPyarrow("pyarrow") + + +def optional_pyarrow( + namespace: Optional[Dict[str, Any]] = None, +) -> Optional[types.ModuleType]: + """``pyarrow``, or None when it isn't installed, without importing it. + + Returns the real module if it is already imported, otherwise a stand-in that + imports it on first attribute access. Whether pyarrow is installed is decided + with ``find_spec``, so ``pyarrow is None`` checks behave as before. Pass the + caller's ``globals()`` to have its ``pyarrow`` replaced by the real module + once the stand-in is first used. + """ + module = sys.modules.get("pyarrow") + if module is not None: + return module + try: + installed = importlib.util.find_spec("pyarrow") is not None + except (ImportError, ValueError): + installed = False + if not installed: + return None + if namespace is not None: + _NAMESPACES.append(namespace) + return _LAZY_PYARROW diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index f24a07505..b5cdfb6f4 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -4,12 +4,13 @@ from typing import List, Optional, TYPE_CHECKING, Tuple import logging -import pandas -try: +from databricks.sql.common.optional_pyarrow import optional_pyarrow + +if TYPE_CHECKING: import pyarrow -except ImportError: - pyarrow = None +else: + pyarrow = optional_pyarrow(globals()) if TYPE_CHECKING: from databricks.sql.backend.thrift_backend import ThriftDatabricksClient @@ -103,6 +104,10 @@ def _convert_arrow_table(self, table): ResultRow(*[v.as_py() for v in r]) for r in zip(*table.itercolumns()) ] + # Imported on first use: pandas takes ~165 ms to import, and connect() + # used to pay that under the import lock for every thread connecting. + import pandas + # Need to use nullable types, as otherwise type can change when there are missing values. # See https://arrow.apache.org/docs/python/pandas.html#nullable-types # NOTE: This api is epxerimental https://pandas.pydata.org/pandas-docs/stable/user_guide/integer_na.html diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 356a96f0f..d59298466 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -25,10 +25,12 @@ import lz4.frame -try: +from databricks.sql.common.optional_pyarrow import optional_pyarrow + +if TYPE_CHECKING: import pyarrow -except ImportError: - pyarrow = None +else: + pyarrow = optional_pyarrow(globals()) from databricks.sql import OperationalError from databricks.sql.exc import ProgrammingError diff --git a/tests/unit/test_lazy_pandas_import.py b/tests/unit/test_lazy_pandas_import.py new file mode 100644 index 000000000..235b43faf --- /dev/null +++ b/tests/unit/test_lazy_pandas_import.py @@ -0,0 +1,113 @@ +"""pandas must load on the first row conversion, not when connect() is imported. + +Each check runs in a fresh interpreter because ``sys.modules`` is process-global. +""" + +import subprocess +import sys +import textwrap + +import pytest + +pytest.importorskip("pyarrow") + +CONNECT_PATH_MODULES = [ + "databricks.sql", + "databricks.sql.client", + "databricks.sql.result_set", + "databricks.sql.utils", + "databricks.sql.backend.thrift_backend", + "databricks.sql.backend.kernel.client", + "databricks.sql.backend.kernel.result_set", + "databricks.sql.backend.sea.backend", +] + + +def _run(script: str) -> subprocess.CompletedProcess: + return subprocess.run( + [sys.executable, "-c", textwrap.dedent(script)], + capture_output=True, + text=True, + timeout=120, + ) + + +def _table_ipc_hex() -> str: + # Children read the table from IPC bytes: building it from Python lists + # makes pyarrow import pandas itself, which would hide what the + # conversion imports. + import io + + import pyarrow + import pyarrow.ipc + + table = pyarrow.table({"a": [1, None], "b": ["x", "y"]}) + sink = io.BytesIO() + with pyarrow.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + return sink.getvalue().hex() + + +def _convert_rows_script(disable_pandas: bool) -> str: + return f""" + import sys + from unittest.mock import MagicMock + import pyarrow.ipc + from databricks.sql.result_set import ResultSet + + table = pyarrow.ipc.open_stream(bytes.fromhex("{_table_ipc_hex()}")).read_all() + assert "pandas" not in sys.modules + rs = MagicMock() + rs.connection.disable_pandas = {disable_pandas} + rs.description = [("a",), ("b",)] + rows = ResultSet._convert_arrow_table(rs, table) + assert [tuple(r) for r in rows] == [(1, "x"), (None, "y")], rows + print("pandas" in sys.modules) + """ + + +# Sentinel exit codes, distinct from the 1 of an uncaught exception. +_NO_PANDAS, _PANDAS_LOADED, _MISSING_OPTIONAL_DEP = 10, 11, 12 +# Only a missing optional extra may skip; any other import error fails. +_OPTIONAL_DEPS = ("databricks_sql_kernel", "pyarrow") + + +@pytest.mark.parametrize("module_name", CONNECT_PATH_MODULES) +def test_connect_path_does_not_import_pandas(module_name): + proc = _run(f""" + import sys + try: + import {module_name} # noqa: F401 + except ImportError as e: + # The kernel backend re-raises a missing wheel as a friendlier + # ImportError, so look for the missing extra along the chain. + cause = e + while cause is not None and not ( + isinstance(cause, ModuleNotFoundError) + and cause.name in {_OPTIONAL_DEPS!r} + ): + cause = cause.__cause__ or cause.__context__ + if cause is None: + raise + print(cause.name) + sys.exit({_MISSING_OPTIONAL_DEP}) + sys.exit({_PANDAS_LOADED} if "pandas" in sys.modules else {_NO_PANDAS}) + """) + if proc.returncode == _MISSING_OPTIONAL_DEP: + pytest.skip(f"{module_name} needs {proc.stdout.strip()}, not installed") + assert proc.returncode != _PANDAS_LOADED, f"importing {module_name} imported pandas" + assert ( + proc.returncode == _NO_PANDAS + ), f"unexpected exit {proc.returncode}:\n{proc.stderr}" + + +def test_row_conversion_imports_pandas_on_first_use(): + proc = _run(_convert_rows_script(disable_pandas=False)) + assert proc.returncode == 0, proc.stderr + assert proc.stdout.strip() == "True", "the pandas path should import pandas" + + +def test_disable_pandas_never_imports_pandas(): + proc = _run(_convert_rows_script(disable_pandas=True)) + assert proc.returncode == 0, proc.stderr + assert proc.stdout.strip() == "False", "_disable_pandas=True imported pandas" diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py new file mode 100644 index 000000000..aa5a29118 --- /dev/null +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -0,0 +1,187 @@ +"""pyarrow (and the numpy it imports) load on the first result, not in connect(). + +Each check runs in a fresh interpreter because ``sys.modules`` is process-global. +""" + +import subprocess +import sys +import textwrap + +import pytest + +from databricks.sql.common import optional_pyarrow as optional_pyarrow_module +from databricks.sql.common.optional_pyarrow import optional_pyarrow + +pytest.importorskip("pyarrow") + +# Sentinel exit codes, distinct from the 1 of an uncaught exception. +_CLEAN, _LOADED = 10, 11 +UNWANTED = ["pyarrow", "numpy", "pandas"] + +# Stand-in for the kernel wheel so a kernel connect() runs without a server. +_FAKE_KERNEL = """ +import sys +import types +from unittest.mock import MagicMock + +fake = types.ModuleType("databricks_sql_kernel") +fake.KernelError = type("KernelError", (Exception,), {}) +fake.Session = MagicMock() +fake.Session.return_value.session_id = "01ef-session" +sys.modules["databricks_sql_kernel"] = fake +""" + + +def _run(body: str) -> subprocess.CompletedProcess: + script = _FAKE_KERNEL + textwrap.dedent(body) + return subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=120 + ) + + +@pytest.mark.parametrize( + "module_name", + [ + "databricks.sql", + "databricks.sql.client", + "databricks.sql.utils", + "databricks.sql.result_set", + "databricks.sql.backend.kernel.client", + ], +) +def test_connect_path_modules_do_not_import_pyarrow(module_name): + # The kernel wheel is faked and pyarrow is present, so any import error is + # a real failure (exit 1), never a skip. + proc = _run(f""" + import {module_name} # noqa: F401 + loaded = [m for m in {UNWANTED!r} if m in sys.modules] + print(loaded) + sys.exit({_LOADED} if loaded else {_CLEAN}) + """) + assert proc.returncode != _LOADED, f"importing {module_name} loaded {proc.stdout}" + assert ( + proc.returncode == _CLEAN + ), f"unexpected exit {proc.returncode}:\n{proc.stderr}" + + +def test_kernel_connect_defers_pyarrow_to_the_first_result(): + proc = _run(f""" + from databricks import sql + + conn = sql.connect( + server_hostname="example.cloud.databricks.com", + http_path="/sql/1.0/warehouses/abc", + access_token="dapi-xyz", + use_kernel=True, + ) + loaded = [m for m in {UNWANTED!r} if m in sys.modules] + assert not loaded, f"connect() loaded {{loaded}}" + handle = MagicMock() + handle.arrow_schema.side_effect = RuntimeError("stop after import") + try: + conn.session.backend._make_result_set(handle, conn.cursor(), None) + except Exception as e: + # Only the sentinel may stop it (the result set wraps it). + cause = e + while cause is not None and str(cause) != "stop after import": + cause = cause.__cause__ or cause.__context__ + if cause is None: + raise + assert "pyarrow" in sys.modules, "the first result imports pyarrow" + """) + assert proc.returncode == 0, proc.stderr + + +def test_without_pyarrow_the_optional_globals_are_none(): + proc = _run(""" + sys.modules["pyarrow"] = None # makes `import pyarrow` fail + import databricks.sql.client as client + import databricks.sql.result_set as result_set + import databricks.sql.utils as utils + + assert client.pyarrow is None + assert utils.pyarrow is None + assert result_set.pyarrow is None + """) + assert proc.returncode == 0, proc.stderr + + +def test_lazy_pyarrow_resolves_to_the_real_module(): + proc = _run(""" + from databricks.sql.common.optional_pyarrow import optional_pyarrow + + lazy = optional_pyarrow() + assert "pyarrow" not in sys.modules and lazy is not None + assert lazy.Table is sys.modules["pyarrow"].Table + import pyarrow + + assert isinstance(pyarrow.table({"a": [1]}), lazy.Table) + assert optional_pyarrow() is pyarrow + """) + assert proc.returncode == 0, proc.stderr + + +def test_first_use_rebinds_the_modules_to_the_real_pyarrow(): + # After the first read, hot-path reads go to the real module, not the + # stand-in's __getattr__. + proc = _run(""" + import databricks.sql.client as client + import databricks.sql.result_set as result_set + import databricks.sql.utils as utils + + stand_in = utils.pyarrow + assert "pyarrow" not in sys.modules + assert all(m.pyarrow is stand_in for m in (client, result_set)) + stand_in.Table + real = sys.modules["pyarrow"] + assert all(m.pyarrow is real for m in (client, utils, result_set)) + assert stand_in.Table is real.Table + """) + assert proc.returncode == 0, proc.stderr + + +def test_rebinding_leaves_a_replaced_global_alone(monkeypatch): + stand_in = optional_pyarrow_module._LazyPyarrow("pyarrow") + replaced, holding = {"pyarrow": "patched"}, {"pyarrow": stand_in} + monkeypatch.setattr(optional_pyarrow_module, "_NAMESPACES", [replaced, holding]) + import pyarrow + + assert stand_in.Table is pyarrow.Table + assert holding["pyarrow"] is pyarrow + assert replaced["pyarrow"] == "patched" + + +def test_stand_in_reads_the_live_module(monkeypatch): + # No copies on the stand-in, so patching pyarrow applies through it too. + import pyarrow + + stand_in = optional_pyarrow_module._LazyPyarrow("pyarrow") + assert stand_in.concat_tables is pyarrow.concat_tables + patched = object() + monkeypatch.setattr(pyarrow, "concat_tables", patched) + assert stand_in.concat_tables is patched + + +def test_optional_pyarrow_is_none_when_not_installed(monkeypatch): + monkeypatch.setitem(sys.modules, "pyarrow", None) + monkeypatch.setattr( + optional_pyarrow_module.importlib.util, "find_spec", lambda name: None + ) + assert optional_pyarrow() is None + + +def test_unimportable_pyarrow_raises_attribute_error(monkeypatch): + attempts = [] + + def broken(name): + attempts.append(name) + raise ImportError("numpy ABI mismatch") + + monkeypatch.setattr(optional_pyarrow_module.importlib, "import_module", broken) + lazy = optional_pyarrow_module._LazyPyarrow("pyarrow") + assert not hasattr(lazy, "Table") + assert getattr(lazy, "Table", None) is None + with pytest.raises(AttributeError, match="numpy ABI mismatch") as raised: + lazy.Table + assert isinstance(raised.value.__cause__, ImportError) + assert attempts == ["pyarrow"], "a failed import is not retried"