From 27479f90d8b707aa9e7bfdab4e5e140e1f65c563 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 18:55:45 +0000 Subject: [PATCH 01/14] perf: import pandas on first row conversion instead of in connect() client.py and result_set.py imported pandas at module load, so the first connect() in a process spent ~165 ms importing it under the import lock while every other connecting thread waited. pandas is only used when an Arrow table is converted to rows, so import it there. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 1 + src/databricks/sql/client.py | 1 - src/databricks/sql/result_set.py | 5 ++- tests/unit/test_lazy_pandas_import.py | 59 +++++++++++++++++++++++++++ 4 files changed, 64 insertions(+), 2 deletions(-) create mode 100644 tests/unit/test_lazy_pandas_import.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 2bae4c54e..9516075e1 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 the first `connect()` in a process by importing pandas only when rows are first converted through it. # 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/client.py b/src/databricks/sql/client.py index 613b62386..eb7516b1e 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -12,7 +12,6 @@ BinaryIO, TYPE_CHECKING, ) -import pandas try: import pyarrow diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index f24a07505..ca039e115 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -4,7 +4,6 @@ from typing import List, Optional, TYPE_CHECKING, Tuple import logging -import pandas try: import pyarrow @@ -103,6 +102,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/tests/unit/test_lazy_pandas_import.py b/tests/unit/test_lazy_pandas_import.py new file mode 100644 index 000000000..60685f59f --- /dev/null +++ b/tests/unit/test_lazy_pandas_import.py @@ -0,0 +1,59 @@ +"""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.client", + "databricks.sql.result_set", + "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, + ) + + +@pytest.mark.parametrize("module_name", CONNECT_PATH_MODULES) +def test_connect_path_does_not_import_pandas(module_name): + proc = _run(f""" + import sys + import {module_name} # noqa: F401 + sys.exit(1 if "pandas" in sys.modules else 0) + """) + if proc.returncode not in (0, 1): + pytest.skip(f"{module_name} not importable here:\n{proc.stderr}") + assert proc.returncode == 0, f"importing {module_name} imported pandas" + + +def test_row_conversion_imports_pandas_on_first_use(): + proc = _run(""" + import sys + from unittest.mock import MagicMock + import pyarrow + from databricks.sql.result_set import ResultSet + + assert "pandas" not in sys.modules + rs = MagicMock() + rs.connection.disable_pandas = False + rs.description = [("a",), ("b",)] + table = pyarrow.table({"a": [1, None], "b": ["x", "y"]}) + rows = ResultSet._convert_arrow_table(rs, table) + assert "pandas" in sys.modules + assert [tuple(r) for r in rows] == [(1, "x"), (None, "y")], rows + """) + assert proc.returncode == 0, proc.stderr From 4041f33dd06c6c3d4e30fca6a4faa4d254197a92 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 19:11:55 +0000 Subject: [PATCH 02/14] test: tell an import failure apart from pandas being loaded Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- tests/unit/test_lazy_pandas_import.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_lazy_pandas_import.py b/tests/unit/test_lazy_pandas_import.py index 60685f59f..9fbaed2ce 100644 --- a/tests/unit/test_lazy_pandas_import.py +++ b/tests/unit/test_lazy_pandas_import.py @@ -28,16 +28,28 @@ def _run(script: str) -> subprocess.CompletedProcess: ) +# Sentinel exit codes, distinct from the 1 of an uncaught exception. +_NO_PANDAS, _PANDAS_LOADED, _IMPORT_FAILED = 10, 11, 12 + + @pytest.mark.parametrize("module_name", CONNECT_PATH_MODULES) def test_connect_path_does_not_import_pandas(module_name): proc = _run(f""" import sys - import {module_name} # noqa: F401 - sys.exit(1 if "pandas" in sys.modules else 0) + try: + import {module_name} # noqa: F401 + except BaseException: + import traceback + traceback.print_exc() + sys.exit({_IMPORT_FAILED}) + sys.exit({_PANDAS_LOADED} if "pandas" in sys.modules else {_NO_PANDAS}) """) - if proc.returncode not in (0, 1): + if proc.returncode == _IMPORT_FAILED: pytest.skip(f"{module_name} not importable here:\n{proc.stderr}") - assert proc.returncode == 0, f"importing {module_name} imported pandas" + 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(): From 9b9111581e0bf235f9cc4678e1c7757f4c6f7633 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 20:32:48 +0000 Subject: [PATCH 03/14] test: skip the pandas import guard only when an optional extra is missing Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- tests/unit/test_lazy_pandas_import.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_lazy_pandas_import.py b/tests/unit/test_lazy_pandas_import.py index 9fbaed2ce..301345a9b 100644 --- a/tests/unit/test_lazy_pandas_import.py +++ b/tests/unit/test_lazy_pandas_import.py @@ -29,7 +29,9 @@ def _run(script: str) -> subprocess.CompletedProcess: # Sentinel exit codes, distinct from the 1 of an uncaught exception. -_NO_PANDAS, _PANDAS_LOADED, _IMPORT_FAILED = 10, 11, 12 +_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) @@ -38,14 +40,15 @@ def test_connect_path_does_not_import_pandas(module_name): import sys try: import {module_name} # noqa: F401 - except BaseException: - import traceback - traceback.print_exc() - sys.exit({_IMPORT_FAILED}) + except ModuleNotFoundError as e: + if e.name not in {_OPTIONAL_DEPS!r}: + raise + print(e.name) + sys.exit({_MISSING_OPTIONAL_DEP}) sys.exit({_PANDAS_LOADED} if "pandas" in sys.modules else {_NO_PANDAS}) """) - if proc.returncode == _IMPORT_FAILED: - pytest.skip(f"{module_name} not importable here:\n{proc.stderr}") + 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 From 1c07208c3d3010c28be7c187b650dd0f79f36a2a Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 20:53:58 +0000 Subject: [PATCH 04/14] test: skip the pandas import guard when the kernel wheel is missing The kernel backend re-raises a missing databricks_sql_kernel as an ImportError, so look for the missing extra along the exception chain. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- tests/unit/test_lazy_pandas_import.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_lazy_pandas_import.py b/tests/unit/test_lazy_pandas_import.py index 301345a9b..1bb0fcfec 100644 --- a/tests/unit/test_lazy_pandas_import.py +++ b/tests/unit/test_lazy_pandas_import.py @@ -40,10 +40,18 @@ def test_connect_path_does_not_import_pandas(module_name): import sys try: import {module_name} # noqa: F401 - except ModuleNotFoundError as e: - if e.name not in {_OPTIONAL_DEPS!r}: + 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(e.name) + print(cause.name) sys.exit({_MISSING_OPTIONAL_DEP}) sys.exit({_PANDAS_LOADED} if "pandas" in sys.modules else {_NO_PANDAS}) """) From 916638bb76e17fa400e70f51c0d194947925989f Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 21:09:41 +0000 Subject: [PATCH 05/14] test: pin that _disable_pandas never imports pandas; bound subprocess time Build the test table from IPC bytes, since constructing it from Python lists makes pyarrow import pandas itself. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- tests/unit/test_lazy_pandas_import.py | 56 ++++++++++++++++++++------- 1 file changed, 42 insertions(+), 14 deletions(-) diff --git a/tests/unit/test_lazy_pandas_import.py b/tests/unit/test_lazy_pandas_import.py index 1bb0fcfec..fc239abaf 100644 --- a/tests/unit/test_lazy_pandas_import.py +++ b/tests/unit/test_lazy_pandas_import.py @@ -25,9 +25,44 @@ def _run(script: str) -> subprocess.CompletedProcess: [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. @@ -64,19 +99,12 @@ def test_connect_path_does_not_import_pandas(module_name): def test_row_conversion_imports_pandas_on_first_use(): - proc = _run(""" - import sys - from unittest.mock import MagicMock - import pyarrow - from databricks.sql.result_set import ResultSet + 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" - assert "pandas" not in sys.modules - rs = MagicMock() - rs.connection.disable_pandas = False - rs.description = [("a",), ("b",)] - table = pyarrow.table({"a": [1, None], "b": ["x", "y"]}) - rows = ResultSet._convert_arrow_table(rs, table) - assert "pandas" in sys.modules - assert [tuple(r) for r in rows] == [(1, "x"), (None, "y")], rows - """) + +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" From c1937110f5370b2a1d22dbceed0293773b304665 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 21:28:39 +0000 Subject: [PATCH 06/14] test: guard the default Thrift connect path against importing pandas Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- tests/unit/test_lazy_pandas_import.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unit/test_lazy_pandas_import.py b/tests/unit/test_lazy_pandas_import.py index fc239abaf..235b43faf 100644 --- a/tests/unit/test_lazy_pandas_import.py +++ b/tests/unit/test_lazy_pandas_import.py @@ -12,8 +12,11 @@ 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", From a2b0a5fc15defbbf7bc8211b72c11e54d3a0bbd1 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 20:25:45 +0000 Subject: [PATCH 07/14] perf(kernel): import pyarrow on the first result instead of in connect() client.py, utils.py and result_set.py imported pyarrow at module load, and the kernel client imported its result set and type mapping (both pyarrow) at load, so a kernel connect() imported pyarrow and numpy. numpy also starts OpenBLAS's thread pool, whose threads busy-wait for a while and take CPU from the connections being opened. The optional pyarrow globals now come from optional_pyarrow(): the real module if it is already imported, None if it isn't installed (decided with find_spec, so pyarrow-is-None checks behave as before), or a stand-in that imports it on first attribute access. The kernel client imports its result set and type mapping where they are first used. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 1 + src/databricks/sql/backend/kernel/client.py | 9 +- src/databricks/sql/client.py | 8 +- src/databricks/sql/common/optional_pyarrow.py | 41 ++++++ src/databricks/sql/result_set.py | 8 +- src/databricks/sql/utils.py | 8 +- tests/unit/test_lazy_pyarrow_import.py | 129 ++++++++++++++++++ 7 files changed, 193 insertions(+), 11 deletions(-) create mode 100644 src/databricks/sql/common/optional_pyarrow.py create mode 100644 tests/unit/test_lazy_pyarrow_import.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9516075e1..bc133f765 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ - 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 eb7516b1e..923e2c31b 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -13,10 +13,12 @@ TYPE_CHECKING, ) -try: +from databricks.sql.common.optional_pyarrow import optional_pyarrow + +if TYPE_CHECKING: import pyarrow -except ImportError: - pyarrow = None +else: + pyarrow = optional_pyarrow() 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..1f378394a --- /dev/null +++ b/src/databricks/sql/common/optional_pyarrow.py @@ -0,0 +1,41 @@ +"""``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. +""" + +import importlib +import importlib.util +import sys +import types +from typing import Any, Optional + + +class _LazyPyarrow(types.ModuleType): + """Stands in for the ``pyarrow`` module until an attribute is used.""" + + def __getattr__(self, name: str) -> Any: + value = getattr(importlib.import_module("pyarrow"), name) + setattr(self, name, value) + return value + + +_LAZY_PYARROW = _LazyPyarrow("pyarrow") + + +def optional_pyarrow() -> 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. + """ + 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 + return _LAZY_PYARROW if installed else None diff --git a/src/databricks/sql/result_set.py b/src/databricks/sql/result_set.py index ca039e115..12b00f2b6 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -5,10 +5,12 @@ import logging -try: +from databricks.sql.common.optional_pyarrow import optional_pyarrow + +if TYPE_CHECKING: import pyarrow -except ImportError: - pyarrow = None +else: + pyarrow = optional_pyarrow() if TYPE_CHECKING: from databricks.sql.backend.thrift_backend import ThriftDatabricksClient diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index 356a96f0f..fbad86472 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() from databricks.sql import OperationalError from databricks.sql.exc import ProgrammingError diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py new file mode 100644 index 000000000..174903ffe --- /dev/null +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -0,0 +1,129 @@ +"""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, _IMPORT_FAILED = 10, 11, 12 +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 + ) + + +@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): + proc = _run(f""" + try: + import {module_name} # noqa: F401 + except BaseException: + import traceback + traceback.print_exc() + sys.exit({_IMPORT_FAILED}) + loaded = [m for m in {UNWANTED!r} if m in sys.modules] + print(loaded) + sys.exit({_LOADED} if loaded else {_CLEAN}) + """) + if proc.returncode == _IMPORT_FAILED: + pytest.skip(f"{module_name} not importable here:\n{proc.stderr}") + 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: + pass + 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_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 From 5216929633dc142660f8fbb103e743ba1d4e500b Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 20:33:19 +0000 Subject: [PATCH 08/14] test: fail, never skip, when a pyarrow import-guard module can't import Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- tests/unit/test_lazy_pyarrow_import.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py index 174903ffe..5115a069d 100644 --- a/tests/unit/test_lazy_pyarrow_import.py +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -15,7 +15,7 @@ pytest.importorskip("pyarrow") # Sentinel exit codes, distinct from the 1 of an uncaught exception. -_CLEAN, _LOADED, _IMPORT_FAILED = 10, 11, 12 +_CLEAN, _LOADED = 10, 11 UNWANTED = ["pyarrow", "numpy", "pandas"] # Stand-in for the kernel wheel so a kernel connect() runs without a server. @@ -50,19 +50,14 @@ def _run(body: str) -> subprocess.CompletedProcess: ], ) 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""" - try: - import {module_name} # noqa: F401 - except BaseException: - import traceback - traceback.print_exc() - sys.exit({_IMPORT_FAILED}) + 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}) """) - if proc.returncode == _IMPORT_FAILED: - pytest.skip(f"{module_name} not importable here:\n{proc.stderr}") assert proc.returncode != _LOADED, f"importing {module_name} loaded {proc.stdout}" assert ( proc.returncode == _CLEAN From aafa88cb3088900eddce2f39c7b70e512236340f Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 20:51:21 +0000 Subject: [PATCH 09/14] fix: raise AttributeError when the lazy pyarrow can't import So hasattr() and getattr(..., default) on the stand-in behave as on a module. Document how a broken pyarrow install differs from before. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- src/databricks/sql/common/optional_pyarrow.py | 12 +++++++++++- tests/unit/test_lazy_pyarrow_import.py | 12 ++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/src/databricks/sql/common/optional_pyarrow.py b/src/databricks/sql/common/optional_pyarrow.py index 1f378394a..6453a1e2d 100644 --- a/src/databricks/sql/common/optional_pyarrow.py +++ b/src/databricks/sql/common/optional_pyarrow.py @@ -3,6 +3,11 @@ ``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 @@ -16,7 +21,12 @@ class _LazyPyarrow(types.ModuleType): """Stands in for the ``pyarrow`` module until an attribute is used.""" def __getattr__(self, name: str) -> Any: - value = getattr(importlib.import_module("pyarrow"), name) + try: + module = importlib.import_module("pyarrow") + except ImportError as e: + # AttributeError, so hasattr() and getattr(..., default) still work. + raise AttributeError(f"pyarrow could not be imported: {e}") from e + value = getattr(module, name) setattr(self, name, value) return value diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py index 5115a069d..8f056bcc3 100644 --- a/tests/unit/test_lazy_pyarrow_import.py +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -122,3 +122,15 @@ def test_optional_pyarrow_is_none_when_not_installed(monkeypatch): optional_pyarrow_module.importlib.util, "find_spec", lambda name: None ) assert optional_pyarrow() is None + + +def test_unimportable_pyarrow_raises_attribute_error(monkeypatch): + def broken(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"): + lazy.Table From 55f717148b08a5dc0b132c3644b8ee72c573e9b1 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 21:10:18 +0000 Subject: [PATCH 10/14] test: bound the pyarrow import-guard subprocesses Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- tests/unit/test_lazy_pyarrow_import.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py index 8f056bcc3..e898bb523 100644 --- a/tests/unit/test_lazy_pyarrow_import.py +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -35,7 +35,7 @@ def _run(body: str) -> subprocess.CompletedProcess: script = _FAKE_KERNEL + textwrap.dedent(body) return subprocess.run( - [sys.executable, "-c", script], capture_output=True, text=True + [sys.executable, "-c", script], capture_output=True, text=True, timeout=120 ) From f77297cbda46372c24a1ebe65e45255570d9562c Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 21:11:33 +0000 Subject: [PATCH 11/14] test: let only the sentinel stop the first-result check Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- tests/unit/test_lazy_pyarrow_import.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py index e898bb523..91cb98e25 100644 --- a/tests/unit/test_lazy_pyarrow_import.py +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -80,8 +80,13 @@ def test_kernel_connect_defers_pyarrow_to_the_first_result(): handle.arrow_schema.side_effect = RuntimeError("stop after import") try: conn.session.backend._make_result_set(handle, conn.cursor(), None) - except Exception: - pass + 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 From 66580dea33c96d73b91dc554a9974bfb1281a25a Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 22:05:47 +0000 Subject: [PATCH 12/14] perf: rebind pyarrow to the real module on first use client, utils and result_set kept the stand-in as their pyarrow global after the first result, so every pyarrow.X on the query path went through a ModuleType subclass's slower attribute lookup. The first use now replaces the stand-in with the real module in each module that registered its globals. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- src/databricks/sql/client.py | 2 +- src/databricks/sql/common/optional_pyarrow.py | 29 ++++++++++++++---- src/databricks/sql/result_set.py | 2 +- src/databricks/sql/utils.py | 2 +- tests/unit/test_lazy_pyarrow_import.py | 30 +++++++++++++++++++ 5 files changed, 57 insertions(+), 8 deletions(-) diff --git a/src/databricks/sql/client.py b/src/databricks/sql/client.py index 923e2c31b..cb7fc0b18 100755 --- a/src/databricks/sql/client.py +++ b/src/databricks/sql/client.py @@ -18,7 +18,7 @@ if TYPE_CHECKING: import pyarrow else: - pyarrow = optional_pyarrow() + 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 index 6453a1e2d..5a052eb34 100644 --- a/src/databricks/sql/common/optional_pyarrow.py +++ b/src/databricks/sql/common/optional_pyarrow.py @@ -14,11 +14,19 @@ import importlib.util import sys import types -from typing import Any, Optional +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.""" + """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. + """ def __getattr__(self, name: str) -> Any: try: @@ -26,6 +34,9 @@ def __getattr__(self, name: str) -> Any: except ImportError as e: # AttributeError, so hasattr() and getattr(..., default) still work. raise AttributeError(f"pyarrow could not be imported: {e}") from e + for namespace in _NAMESPACES: + if namespace.get("pyarrow") is self: + namespace["pyarrow"] = module value = getattr(module, name) setattr(self, name, value) return value @@ -34,12 +45,16 @@ def __getattr__(self, name: str) -> Any: _LAZY_PYARROW = _LazyPyarrow("pyarrow") -def optional_pyarrow() -> Optional[types.ModuleType]: +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. + 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: @@ -48,4 +63,8 @@ def optional_pyarrow() -> Optional[types.ModuleType]: installed = importlib.util.find_spec("pyarrow") is not None except (ImportError, ValueError): installed = False - return _LAZY_PYARROW if installed else None + 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 12b00f2b6..b5cdfb6f4 100644 --- a/src/databricks/sql/result_set.py +++ b/src/databricks/sql/result_set.py @@ -10,7 +10,7 @@ if TYPE_CHECKING: import pyarrow else: - pyarrow = optional_pyarrow() + pyarrow = optional_pyarrow(globals()) if TYPE_CHECKING: from databricks.sql.backend.thrift_backend import ThriftDatabricksClient diff --git a/src/databricks/sql/utils.py b/src/databricks/sql/utils.py index fbad86472..d59298466 100644 --- a/src/databricks/sql/utils.py +++ b/src/databricks/sql/utils.py @@ -30,7 +30,7 @@ if TYPE_CHECKING: import pyarrow else: - pyarrow = optional_pyarrow() + pyarrow = optional_pyarrow(globals()) from databricks.sql import OperationalError from databricks.sql.exc import ProgrammingError diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py index 91cb98e25..88bd1f37c 100644 --- a/tests/unit/test_lazy_pyarrow_import.py +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -121,6 +121,36 @@ def test_lazy_pyarrow_resolves_to_the_real_module(): 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_optional_pyarrow_is_none_when_not_installed(monkeypatch): monkeypatch.setitem(sys.modules, "pyarrow", None) monkeypatch.setattr( From 517c4c8d8b25c60ff1017de7579b3c267d53f817 Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 22:31:29 +0000 Subject: [PATCH 13/14] fix: try a failed pyarrow import only once When pyarrow is installed but can't be imported, every attribute read on the stand-in retried the import, rerunning pyarrow's and numpy's partial initialization. The first ImportError is now kept and later reads raise from it, as the eager import ran it once. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- src/databricks/sql/common/optional_pyarrow.py | 14 ++++++++++---- tests/unit/test_lazy_pyarrow_import.py | 7 ++++++- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/common/optional_pyarrow.py b/src/databricks/sql/common/optional_pyarrow.py index 5a052eb34..42291eecc 100644 --- a/src/databricks/sql/common/optional_pyarrow.py +++ b/src/databricks/sql/common/optional_pyarrow.py @@ -29,11 +29,17 @@ class _LazyPyarrow(types.ModuleType): """ def __getattr__(self, name: str) -> Any: - try: - module = importlib.import_module("pyarrow") - except ImportError as e: + # 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: {e}") from e + raise AttributeError(f"pyarrow could not be imported: {error}") from error for namespace in _NAMESPACES: if namespace.get("pyarrow") is self: namespace["pyarrow"] = module diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py index 88bd1f37c..30d240fd3 100644 --- a/tests/unit/test_lazy_pyarrow_import.py +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -160,12 +160,17 @@ def test_optional_pyarrow_is_none_when_not_installed(monkeypatch): 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"): + 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" From 9770a691eeaf0dff56dae47ad8fdb0585220822f Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Sat, 14 Nov 2026 00:23:32 +0000 Subject: [PATCH 14/14] fix: read pyarrow attributes from the live module through the stand-in The stand-in copied each resolved attribute onto itself, so a later patch of pyarrow (mock.patch, monkeypatch, shims) didn't reach code reading through it. It now keeps no copies and always delegates to the imported module; the hot path reads the real module directly since the rebinding. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- src/databricks/sql/common/optional_pyarrow.py | 7 +++---- tests/unit/test_lazy_pyarrow_import.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/databricks/sql/common/optional_pyarrow.py b/src/databricks/sql/common/optional_pyarrow.py index 42291eecc..b67d28b5a 100644 --- a/src/databricks/sql/common/optional_pyarrow.py +++ b/src/databricks/sql/common/optional_pyarrow.py @@ -25,7 +25,8 @@ class _LazyPyarrow(types.ModuleType): 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. + 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: @@ -43,9 +44,7 @@ def __getattr__(self, name: str) -> Any: for namespace in _NAMESPACES: if namespace.get("pyarrow") is self: namespace["pyarrow"] = module - value = getattr(module, name) - setattr(self, name, value) - return value + return getattr(module, name) _LAZY_PYARROW = _LazyPyarrow("pyarrow") diff --git a/tests/unit/test_lazy_pyarrow_import.py b/tests/unit/test_lazy_pyarrow_import.py index 30d240fd3..aa5a29118 100644 --- a/tests/unit/test_lazy_pyarrow_import.py +++ b/tests/unit/test_lazy_pyarrow_import.py @@ -151,6 +151,17 @@ def test_rebinding_leaves_a_replaced_global_alone(monkeypatch): 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(