perf(kernel): import pyarrow on the first result instead of in connect() (stacked on #1049)

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.

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,

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

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

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

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

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"