lazy load pandas and numpy to improve startup performance (#443)

Author: vangheem

CHANGELOG.md

@@ -1,3 +1,9 @@
+## 0.5.21
+
+### Fixes
+
+* **Lazy load pandas and numpy** to improve startup performance
+
 ## 0.5.20
 
 ### Features 

requirements/common/base.in

@@ -1,7 +1,6 @@
 -c constraints.txt
 
 python-dateutil
-pandas
 # Pydantic generic Secret only introduced in 2.7
 pydantic>=2.7
 dataclasses_json

requirements/common/extras.in

@@ -0,0 +1,5 @@
+# Requirements file that is only used by extras
+# These requirements will not be installed by default(`unstructured-ingest`)
+# but will be installed when `unstructured-ingest[<any extra>]` is used
+pandas
+numpy

setup.py

@@ -25,6 +25,8 @@
 
 from unstructured_ingest.__version__ import __version__
 
+_extra_reqs_filepath = "requirements/common/extras.in"
+
 
 def load_requirements(file: Union[str, Path]) -> List[str]:
     path = file if isinstance(file, Path) else Path(file)
@@ -42,8 +44,13 @@ def load_requirements(file: Union[str, Path]) -> List[str]:
         file_spec = recursive_req.split()[-1]
         file_path = Path(file_dir) / file_spec
         requirements.extend(load_requirements(file=file_path.resolve()))
+
     # Remove duplicates and any blank entries
-    return list({r for r in requirements if r})
+    result = list({r for r in requirements if r})
+
+    if file != _extra_reqs_filepath:
+        result.extend(load_requirements(_extra_reqs_filepath))
+    return result
 
 
 csv_reqs = load_requirements("requirements/local_partition/tsv.in")

test/integration/connectors/test_astradb.py

@@ -1,3 +1,4 @@
+import contextlib
 import json
 import os
 from dataclasses import dataclass
@@ -231,6 +232,13 @@ def test_astra_create_destination():
     )
     collection_name = "system_created-123"
     formatted_collection_name = "system_created_123"
+
+    client = AstraDBClient()
+    db = client.get_database(api_endpoint=env_data.api_endpoint, token=env_data.token)
+    with contextlib.suppress(Exception):
+        # drop collection before trying to create it
+        db.drop_collection(formatted_collection_name)
+
     created = uploader.create_destination(destination_name=collection_name, vector_length=3072)
     assert created
     assert uploader.upload_config.collection_name == formatted_collection_name
@@ -239,8 +247,6 @@ def test_astra_create_destination():
     assert not created
 
     # cleanup
-    client = AstraDBClient()
-    db = client.get_database(api_endpoint=env_data.api_endpoint, token=env_data.token)
     db.drop_collection(formatted_collection_name)
 
 

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.5.20"  # pragma: no cover
+__version__ = "0.5.21"  # pragma: no cover

unstructured_ingest/embed/interfaces.py

@@ -2,10 +2,10 @@
 from dataclasses import dataclass
 from typing import Any, Optional
 
-import numpy as np
 from pydantic import BaseModel, Field
 
 from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.utils.dep_check import requires_dependencies
 
 EMBEDDINGS_KEY = "embeddings"
 
@@ -32,7 +32,6 @@ def wrap_error(self, e: Exception) -> Exception:
 
 @dataclass
 class BaseEmbeddingEncoder(BaseEncoder, ABC):
-
     def initialize(self):
         """Initializes the embedding encoder class. Should also validate the instance
         is properly configured: e.g., embed a single a element"""
@@ -46,8 +45,11 @@ def get_exemplary_embedding(self) -> list[float]:
         return self.embed_query(query="Q")
 
     @property
+    @requires_dependencies(["numpy"])
     def is_unit_vector(self) -> bool:
         """Denotes if the embedding vector is a unit vector."""
+        import numpy as np
+
         exemplary_embedding = self.get_exemplary_embedding()
         return np.isclose(np.linalg.norm(exemplary_embedding), 1.0, rtol=1e-03)
 
@@ -86,7 +88,6 @@ def embed_query(self, query: str) -> list[float]:
 
 @dataclass
 class AsyncBaseEmbeddingEncoder(BaseEncoder, ABC):
-
     async def initialize(self):
         """Initializes the embedding encoder class. Should also validate the instance
         is properly configured: e.g., embed a single a element"""
@@ -100,8 +101,11 @@ async def get_exemplary_embedding(self) -> list[float]:
         return await self.embed_query(query="Q")
 
     @property
+    @requires_dependencies(["numpy"])
     async def is_unit_vector(self) -> bool:
         """Denotes if the embedding vector is a unit vector."""
+        import numpy as np
+
         exemplary_embedding = await self.get_exemplary_embedding()
         return np.isclose(np.linalg.norm(exemplary_embedding), 1.0, rtol=1e-03)
 

unstructured_ingest/utils/data_prep.py

@@ -2,20 +2,22 @@
 import json
 from datetime import datetime
 from pathlib import Path
-from typing import Any, Generator, Iterable, Optional, Sequence, TypeVar, Union, cast
-
-import pandas as pd
+from typing import TYPE_CHECKING, Any, Generator, Iterable, Optional, Sequence, TypeVar, Union, cast
 
 from unstructured_ingest.utils import ndjson
+from unstructured_ingest.utils.dep_check import requires_dependencies
 from unstructured_ingest.v2.logger import logger
 
+if TYPE_CHECKING:
+    from pandas import DataFrame
+
 DATE_FORMATS = ("%Y-%m-%d", "%Y-%m-%dT%H:%M:%S", "%Y-%m-%d+%H:%M:%S", "%Y-%m-%dT%H:%M:%S%z")
 
 T = TypeVar("T")
 IterableT = Iterable[T]
 
 
-def split_dataframe(df: pd.DataFrame, chunk_size: int = 100) -> Generator[pd.DataFrame, None, None]:
+def split_dataframe(df: "DataFrame", chunk_size: int = 100) -> Generator["DataFrame", None, None]:
     num_chunks = len(df) // chunk_size + 1
     for i in range(num_chunks):
         yield df[i * chunk_size : (i + 1) * chunk_size]
@@ -144,9 +146,13 @@ def get_data_by_suffix(path: Path) -> list[dict]:
         elif path.suffix == ".ndjson":
             return ndjson.load(f)
         elif path.suffix == ".csv":
+            import pandas as pd
+
             df = pd.read_csv(path)
             return df.to_dict(orient="records")
         elif path.suffix == ".parquet":
+            import pandas as pd
+
             df = pd.read_parquet(path)
             return df.to_dict(orient="records")
         else:
@@ -180,6 +186,9 @@ def get_data(path: Union[Path, str]) -> list[dict]:
             return ndjson.load(f)
         except Exception as e:
             logger.warning(f"failed to read {path} as ndjson: {e}")
+
+        import pandas as pd
+
         try:
             df = pd.read_csv(path)
             return df.to_dict(orient="records")
@@ -202,7 +211,10 @@ def get_json_data(path: Path) -> list[dict]:
             raise ValueError(f"Unsupported file type: {path}")
 
 
-def get_data_df(path: Path) -> pd.DataFrame:
+@requires_dependencies(["pandas"])
+def get_data_df(path: Path) -> "DataFrame":
+    import pandas as pd
+
     with path.open() as f:
         if path.suffix == ".json":
             data = json.load(f)

unstructured_ingest/utils/table.py

@@ -1,11 +1,16 @@
-from typing import Any
-
-import pandas as pd
+from typing import TYPE_CHECKING, Any
 
 from unstructured_ingest.utils.data_prep import flatten_dict
+from unstructured_ingest.utils.dep_check import requires_dependencies
+
+if TYPE_CHECKING:
+    from pandas import DataFrame
 
 
+@requires_dependencies(["pandas"])
 def get_default_pandas_dtypes() -> dict[str, Any]:
+    import pandas as pd
+
     return {
         "text": pd.StringDtype(),  # type: ignore
         "type": pd.StringDtype(),  # type: ignore
@@ -57,7 +62,9 @@ def get_default_pandas_dtypes() -> dict[str, Any]:
 def convert_to_pandas_dataframe(
     elements_dict: list[dict[str, Any]],
     drop_empty_cols: bool = False,
-) -> pd.DataFrame:
+) -> "DataFrame":
+    import pandas as pd
+
     # Flatten metadata if it hasn't already been flattened
     for d in elements_dict:
         if metadata := d.pop("metadata", None):

unstructured_ingest/v2/processes/connectors/delta_table.py

@@ -3,10 +3,9 @@
 from dataclasses import dataclass, field
 from multiprocessing import Process, Queue
 from pathlib import Path
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 from urllib.parse import urlparse
 
-import pandas as pd
 from pydantic import Field, Secret
 
 from unstructured_ingest.error import DestinationConnectionError
@@ -27,6 +26,9 @@
 
 CONNECTOR_TYPE = "delta_table"
 
+if TYPE_CHECKING:
+    from pandas import DataFrame
+
 
 @requires_dependencies(["deltalake"], extras="delta-table")
 def write_deltalake_with_error_handling(queue, **kwargs):
@@ -136,7 +138,7 @@ def precheck(self):
                 logger.error(f"failed to validate connection: {e}", exc_info=True)
                 raise DestinationConnectionError(f"failed to validate connection: {e}")
 
-    def upload_dataframe(self, df: pd.DataFrame, file_data: FileData) -> None:
+    def upload_dataframe(self, df: "DataFrame", file_data: FileData) -> None:
         updated_upload_path = os.path.join(
             self.connection_config.table_uri, file_data.source_identifiers.relative_path
         )
@@ -172,7 +174,10 @@ def upload_dataframe(self, df: pd.DataFrame, file_data: FileData) -> None:
             logger.error(f"Exception occurred in write_deltalake: {error_message}")
             raise RuntimeError(f"Error in write_deltalake: {error_message}")
 
+    @requires_dependencies(["pandas"], extras="delta-table")
     def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        import pandas as pd
+
         df = pd.DataFrame(data=data)
         self.upload_dataframe(df=df, file_data=file_data)