Make imports of heavy packages lazy (#491)

* Lazy imports to make pandas, pyarrow, polars, and duckdb not import on VegaFusion import

* Update package.jsons

* Add lazy import test

* Revert "Update package.jsons"

This reverts commit afd2f71.

* Add altair to lazy import checks

* return it 🤦

Author: jonmmease

.github/workflows/build_test.yml

@@ -320,6 +320,9 @@ jobs:
           python -m pip install vegafusion-*.whl
           python -m pip install vegafusion_python_embed-*manylinux_2_17_x86_64*.whl
           python -m pip install pytest vega-datasets polars-lts-cpu duckdb==0.9.2 "vl-convert-python>=1.0.1rc1" scikit-image pandas==2.0
+      - name: Test lazy imports
+        working-directory: python/vegafusion/
+        run: python checks/check_lazy_imports.py
       - name: Test vegafusion
         working-directory: python/vegafusion/
         run: pytest

python/vegafusion/checks/README.md

@@ -0,0 +1,4 @@
+## Checks
+Python scripts to run on CI that should not be run as part of the pytest test suite.
+In particular, to check lazy imports.
+

python/vegafusion/checks/check_lazy_imports.py

@@ -0,0 +1,6 @@
+import sys
+import vegafusion as vf
+import vegafusion_embed
+
+for mod in ["polars", "pandas", "pyarrow", "duckdb", "altair"]:
+    assert mod not in sys.modules, f"{mod} module should be imported lazily"

python/vegafusion/setup.cfg

@@ -27,7 +27,7 @@ packages = find:
 python_requires = >=3.7
 include_package_data = True
 install_requires = 
-	altair>=4.2.0
+	altair>=5.2.0
 	pyarrow>=5
 	pandas
 	psutil
@@ -38,3 +38,10 @@ embed =
 	vegafusion-python-embed==1.6.8
 	vl-convert-python>=0.7.0
 
+[options.entry_points]
+altair.vegalite.v5.renderer =
+    vegafusion-mime = vegafusion.renderer:vegafusion_mime_renderer
+
+altair.vegalite.v5.data_transformer =
+    vegafusion-feather = vegafusion.transformer:feather_transformer
+    vegafusion-inline = vegafusion.transformer:inline_data_transformer

python/vegafusion/vegafusion/__init__.py

@@ -5,8 +5,7 @@
 from .evaluation import transformed_data
 from .save import save_html, save_vega, save_png, save_svg
 from . import renderer
-from .compilers import vegalite_compilers
-import altair as alt
+
 
 from ._version import __version__
 from .utils import RendererTransformerEnabler
@@ -59,6 +58,7 @@ def enable(mimetype="html", row_limit=10000, embed_options=None):
         Dictionary of options to pass to the vega-embed. Default
         entry is {'mode': 'vega'}.
     """
+    import altair as alt
     embed_options = embed_options if embed_options is not None else {}
     return RendererTransformerEnabler(
         renderer_ctx=alt.renderers.enable(
@@ -115,6 +115,7 @@ def disable():
 
     This does not affect the behavior of VegaFusionWidget
     """
+    import altair as alt
     return RendererTransformerEnabler(
         renderer_ctx=alt.renderers.enable('default'),
         data_transformer_ctx=alt.data_transformers.enable('default'),

python/vegafusion/vegafusion/compilers.py

@@ -1,11 +1,11 @@
-from typing import Dict, Optional
-
-import pandas as pd
-import pyarrow as pa
-
+from typing import Dict, Optional, TYPE_CHECKING
 from dataclasses import dataclass
 from abc import ABC, abstractmethod
 
+if TYPE_CHECKING:
+    from pyarrow import Schema, Table
+    from pandas import DataFrame
+
 
 @dataclass
 class CsvReadOptions:
@@ -15,7 +15,7 @@ class CsvReadOptions:
     has_header: bool
     delimeter: str
     file_extension: str
-    schema: Optional[pa.Schema]
+    schema: Optional["Schema"]
 
 
 class RegistrationNotSupportedError(RuntimeError):
@@ -46,7 +46,7 @@ def dialect(cls) -> str:
         raise NotImplementedError()
 
     @abstractmethod
-    def tables(self) -> Dict[str, pa.Schema]:
+    def tables(self) -> Dict[str, "Schema"]:
         """
         Returns the names and schema for the tables that are provided by the connection.
         These are the tables that may be referenced by SQL queries passed to the
@@ -57,7 +57,7 @@ def tables(self) -> Dict[str, pa.Schema]:
         raise NotImplementedError()
 
     @abstractmethod
-    def fetch_query(self, query: str, schema: pa.Schema) -> pa.Table:
+    def fetch_query(self, query: str, schema: "Schema") -> "Table":
         """
         Returns the result of evaluating the requested query. The resulting pa.Table
         should have a schema matching the provided schema
@@ -77,7 +77,7 @@ def fallback(self) -> bool:
         """
         return True
 
-    def register_pandas(self, name: str, df: pd.DataFrame, temporary: bool = False):
+    def register_pandas(self, name: str, df: "DataFrame", temporary: bool = False):
         """
         Register the provided pandas DataFrame as a table with the provided name
 
@@ -88,7 +88,7 @@ def register_pandas(self, name: str, df: pd.DataFrame, temporary: bool = False):
         """
         raise RegistrationNotSupportedError("Connection does not support registration of pandas datasets")
 
-    def register_arrow(self, name: str, table: pa.Table, temporary: bool = False):
+    def register_arrow(self, name: str, table: "Table", temporary: bool = False):
         """
         Register the provided pyarrow Table as a table with the provided name
         :param name: Table name

python/vegafusion/vegafusion/connection/__init__.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
-import pyarrow as pa
-from vegafusion.proto.datafusion_pb2 import LogicalExprNode
+from typing import Optional, List, Literal, Union, Any, TYPE_CHECKING
 
-from typing import Optional, List, Literal, Union, Any
+if TYPE_CHECKING:
+    from pyarrow import Schema, Table
+    from vegafusion.proto.datafusion_pb2 import LogicalExprNode
 
 
 class DataFrameOperationNotSupportedError(RuntimeError):
@@ -19,12 +20,12 @@ class DataFrameDataset(ABC):
     """
 
     @abstractmethod
-    def schema(self) -> pa.Schema:
+    def schema(self) -> "Schema":
         """DataFrame's pyarrow schema"""
         raise NotImplementedError()
 
     @abstractmethod
-    def collect(self) -> pa.Table:
+    def collect(self) -> "Table":
         """Return DataFrame's value as a pyarrow Table"""
         raise NotImplementedError()
 
@@ -47,7 +48,7 @@ def main_thread(self) -> bool:
         return True
 
     def sort(
-        self, exprs: List[LogicalExprNode], limit: Optional[int]
+        self, exprs: List["LogicalExprNode"], limit: Optional[int]
     ) -> "DataFrameDataset":
         """
         Sort and optionally limit dataset
@@ -58,7 +59,7 @@ def sort(
         """
         raise DataFrameOperationNotSupportedError()
 
-    def select(self, exprs: List[LogicalExprNode]) -> "DataFrameDataset":
+    def select(self, exprs: List["LogicalExprNode"]) -> "DataFrameDataset":
         """
         Select columns from Dataset. Selection expressions may include column names,
         column expressions, or window expressions
@@ -69,7 +70,7 @@ def select(self, exprs: List[LogicalExprNode]) -> "DataFrameDataset":
         raise DataFrameOperationNotSupportedError()
 
     def aggregate(
-        self, group_exprs: List[LogicalExprNode], agg_exprs: List[LogicalExprNode]
+        self, group_exprs: List["LogicalExprNode"], agg_exprs: List["LogicalExprNode"]
     ) -> "DataFrameDataset":
         """
         Perform dataset aggregation. Resulting dataset includes grouping
@@ -82,7 +83,7 @@ def aggregate(
         raise DataFrameOperationNotSupportedError()
 
     def joinaggregate(
-        self, group_exprs: List[LogicalExprNode], agg_exprs: List[LogicalExprNode]
+        self, group_exprs: List["LogicalExprNode"], agg_exprs: List["LogicalExprNode"]
     ) -> "DataFrameDataset":
         """
         Perform joinaggregate dataset operation.
@@ -95,7 +96,7 @@ def joinaggregate(
         """
         raise DataFrameOperationNotSupportedError()
 
-    def filter(self, predicate: LogicalExprNode) -> "DataFrameDataset":
+    def filter(self, predicate: "LogicalExprNode") -> "DataFrameDataset":
         """
         Filter dataset by predicate expression
 
@@ -135,7 +136,7 @@ def fold(
     def stack(
         self,
         field: str,
-        orderby: List[LogicalExprNode],
+        orderby: List["LogicalExprNode"],
         groupby: List[str],
         start_field: str,
         stop_field: str,

python/vegafusion/vegafusion/dataset/dataframe.py

@@ -1,8 +1,9 @@
-import pyarrow as pa
 from functools import cached_property
-from typing import Sequence, Optional
+from typing import Sequence, Optional, TYPE_CHECKING
 from .sql import SqlDataset
 
+if TYPE_CHECKING:
+    from pyarrow import Table, Schema
 
 class SqlDatasetDataFrame:
     """An implementation of the dataframe interchange protocol.
@@ -22,7 +23,7 @@ def __init__(
         dataset: SqlDataset,
         nan_as_null: bool = False,
         allow_copy: bool = True,
-        pyarrow_table: Optional[pa.Table] = None,
+        pyarrow_table: Optional["Table"] = None,
     ):
         self._dataset = dataset
         self._nan_as_null = nan_as_null
@@ -55,7 +56,7 @@ def _empty_pyarrow_df(self):
         return schema.empty_table().__dataframe__()
 
     @property
-    def _schema(self) -> pa.Schema:
+    def _schema(self) -> "Schema":
         return self._dataset.table_schema()
 
     def _get_dtype(self, name):

python/vegafusion/vegafusion/dataset/dfi.py

@@ -1,9 +1,9 @@
 from abc import ABC, abstractmethod
-import pyarrow as pa
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from .dfi import SqlDatasetDataFrame
+    import pyarrow as pa
 
 
 class SqlDataset(ABC):
@@ -35,12 +35,12 @@ def table_name(self) -> str:
         raise NotImplementedError()
 
     @abstractmethod
-    def table_schema(self) -> pa.Schema:
+    def table_schema(self) -> "pa.Schema":
         """Schema of source table"""
         raise NotImplementedError()
 
     @abstractmethod
-    def fetch_query(self, query: str, schema: pa.Schema) -> pa.Table:
+    def fetch_query(self, query: str, schema: "pa.Schema") -> "pa.Table":
         """
         Returns the result of evaluating the requested query. The resulting pa.Table
         should have a schema matching the provided schema