feat: implement lazy loading for pyarrow to optimize imports

Author: kosiew

python/datafusion/expr.py

@@ -40,6 +40,18 @@
 if TYPE_CHECKING:
     from datafusion.plan import LogicalPlan
 
+
+class _LazyPyArrow:
+    """Lazily import :mod:`pyarrow` on first attribute access."""
+
+    def __getattr__(self, name: str) -> Any:  # pragma: no cover - optional dependency
+        import pyarrow as pa  # type: ignore[import]
+
+        return getattr(pa, name)
+
+
+pa = _LazyPyArrow()
+
 # The following are imported from the internal representation. We may choose to
 # give these all proper wrappers, or to simply leave as is. These were added
 # in order to support passing the `test_imports` unit test.

python/datafusion/functions.py

@@ -20,8 +20,6 @@
 
 from typing import TYPE_CHECKING, Any, Optional
 
-import pyarrow as pa
-
 from datafusion._internal import functions as f
 from datafusion.common import NullTreatment
 from datafusion.expr import (
@@ -36,6 +34,18 @@
 if TYPE_CHECKING:
     from datafusion.context import SessionContext
 
+
+class _LazyPyArrow:
+    """Lazily import :mod:`pyarrow` on first attribute access."""
+
+    def __getattr__(self, name: str) -> Any:  # pragma: no cover - optional dependency
+        import pyarrow as pa  # type: ignore[import]
+
+        return getattr(pa, name)
+
+
+pa = _LazyPyArrow()
+
 __all__ = [
     "abs",
     "acos",

python/datafusion/user_defined.py

@@ -24,13 +24,27 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, TypeVar, overload
 
-import pyarrow as pa
-
 import datafusion._internal as df_internal
 from datafusion.expr import Expr
 
 if TYPE_CHECKING:
+    import pyarrow as pa
+
     _R = TypeVar("_R", bound=pa.DataType)
+else:
+    _R = TypeVar("_R")
+
+
+class _LazyPyArrow:
+    """Lazily import :mod:`pyarrow` on first attribute access."""
+
+    def __getattr__(self, name: str) -> Any:  # pragma: no cover - optional dependency
+        import pyarrow as pa  # type: ignore[import]
+
+        return getattr(pa, name)
+
+
+pa = _LazyPyArrow()
 
 
 class Volatility(Enum):