Fix arrow_flight_client imports (#329)

* Fix arrow_flight_client imports

* Provide human-readable errors in case Pyarrow is not installed

* Ruff

Author: aversey

python/hopsworks_common/core/constants.py

@@ -25,25 +25,35 @@
 HAS_CONFLUENT_KAFKA: bool = importlib.util.find_spec("confluent_kafka") is not None
 confluent_kafka_not_installed_message = (
     "Confluent Kafka package not found. "
-    "If you want to use Kafka with Hopsworks you can install the corresponding extras "
-    """`pip install hopsworks[python]` or `pip install "hopsworks[python]"` if using zsh. """
-    "You can also install confluent-kafka directly in your environment e.g `pip install confluent-kafka`. "
+    "If you want to use Kafka with Hopsworks you can install the corresponding extras via "
+    '`pip install "hopsworks[python]"`. '
+    "You can also install confluent-kafka directly in your environment with `pip install confluent-kafka`. "
     "You will need to restart your kernel if applicable."
 )
+
 # Data Validation / Great Expectations
 HAS_GREAT_EXPECTATIONS: bool = (
     importlib.util.find_spec("great_expectations") is not None
 )
 great_expectations_not_installed_message = (
     "Great Expectations package not found. "
-    "If you want to use data validation with Hopsworks you can install the corresponding extras "
-    """`pip install hopsworks[great_expectations]` or `pip install "hopsworks[great_expectations]"` if using zsh. """
-    "You can also install great-expectations directly in your environment e.g `pip install great-expectations`. "
+    "If you want to use data validation with Hopsworks you can install the corresponding extras via "
+    '`pip install "hopsworks[great_expectations]"`. '
+    "You can also install great-expectations directly in your environment with `pip install great-expectations`. "
     "You will need to restart your kernel if applicable."
 )
 initialise_expectation_suite_for_single_expectation_api_message = "Initialize Expectation Suite by attaching to a Feature Group to enable single expectation API"
 
-HAS_ARROW: bool = importlib.util.find_spec("pyarrow") is not None
+# Pyarrow
+HAS_PYARROW: bool = importlib.util.find_spec("pyarrow") is not None
+pyarrow_not_installed_message = (
+    "Pyarrow package not found. "
+    "If you want to use Apache Arrow with Hopsworks you can install the corresponding extras via "
+    '`pip install "hopsworks[python]"`. '
+    "You can also install pyarrow directly in your environment with `pip install pyarrow`. "
+    "You will need to restart your kernel if applicable."
+)
+
 HAS_PANDAS: bool = importlib.util.find_spec("pandas") is not None
 HAS_NUMPY: bool = importlib.util.find_spec("numpy") is not None
 HAS_POLARS: bool = importlib.util.find_spec("polars") is not None

python/hopsworks_common/core/type_systems.py

@@ -23,9 +23,9 @@
 
 import pytz
 from hopsworks_common.core.constants import (
-    HAS_ARROW,
     HAS_PANDAS,
     HAS_POLARS,
+    HAS_PYARROW,
 )
 from hopsworks_common.decorators import uses_polars
 
@@ -35,7 +35,7 @@
     import pandas as pd
     import polars as pl
 
-if HAS_ARROW:
+if HAS_PYARROW:
     import pyarrow as pa
 
     # Decimal types are currently not supported

python/hsfs/constructor/query.py

@@ -29,7 +29,7 @@
 from hsfs.constructor import join
 from hsfs.constructor.filter import Filter, Logic
 from hsfs.constructor.fs_query import FsQuery
-from hsfs.core import arrow_flight_client, query_constructor_api, storage_connector_api
+from hsfs.core import query_constructor_api, storage_connector_api
 from hsfs.decorators import typechecked
 from hsfs.feature import Feature
 
@@ -101,6 +101,8 @@ def _prep_read(
             online_conn = None
 
             if engine.get_instance().is_flyingduck_query_supported(self, read_options):
+                from hsfs.core import arrow_flight_client
+
                 sql_query = self._to_string(fs_query, online, asof=True)
                 sql_query = arrow_flight_client.get_instance().create_query_object(
                     self, sql_query, fs_query.on_demand_fg_aliases

python/hsfs/core/arrow_flight_client.py

@@ -23,6 +23,12 @@
 from functools import wraps
 from typing import Any, Dict, Optional, Union
 
+from hopsworks_common.core.constants import HAS_PYARROW, pyarrow_not_installed_message
+
+
+if not HAS_PYARROW:
+    raise ModuleNotFoundError(pyarrow_not_installed_message)
+
 import pyarrow
 import pyarrow._flight
 import pyarrow.flight

python/hsfs/core/constants.py

@@ -16,14 +16,14 @@
 
 from hopsworks_common.core.constants import (
     HAS_AIOMYSQL,
-    HAS_ARROW,
     HAS_AVRO,
     HAS_CONFLUENT_KAFKA,
     HAS_FAST_AVRO,
     HAS_GREAT_EXPECTATIONS,
     HAS_NUMPY,
     HAS_PANDAS,
     HAS_POLARS,
+    HAS_PYARROW,
     HAS_SQLALCHEMY,
     great_expectations_not_installed_message,
     initialise_expectation_suite_for_single_expectation_api_message,
@@ -32,7 +32,7 @@
 
 __all__ = [
     "HAS_AIOMYSQL",
-    "HAS_ARROW",
+    "HAS_PYARROW",
     "HAS_AVRO",
     "HAS_CONFLUENT_KAFKA",
     "HAS_FAST_AVRO",

python/hsfs/core/feature_view_engine.py

@@ -33,7 +33,6 @@
 from hsfs.client import exceptions
 from hsfs.constructor.filter import Filter, Logic
 from hsfs.core import (
-    arrow_flight_client,
     feature_view_api,
     query_constructor_api,
     statistics_engine,
@@ -1009,7 +1008,17 @@ def get_models_provenance(
 
     def _check_feature_group_accessibility(self, feature_view_obj):
         if engine.get_type() == "python":
-            if arrow_flight_client.get_instance().is_enabled():
+            try:
+                from hsfs.core import arrow_flight_client
+
+                arrow_flight_client_imported = True
+            except ImportError:
+                arrow_flight_client_imported = False
+
+            if (
+                arrow_flight_client_imported
+                and arrow_flight_client.get_instance().is_enabled()
+            ):
                 if not arrow_flight_client.supports(
                     feature_view_obj.query.featuregroups
                 ):
@@ -1227,7 +1236,9 @@ def _get_feature_logging_data(
                 model_col_name=FeatureViewEngine._HSML_MODEL,
                 predictions=predictions,
                 training_dataset_version=training_dataset_version,
-                hsml_model=self.get_hsml_model_value(hsml_model) if hsml_model else None,
+                hsml_model=self.get_hsml_model_value(hsml_model)
+                if hsml_model
+                else None,
             )
         else:
             return engine.get_instance().get_feature_logging_df(

python/hsfs/engine/__init__.py

@@ -19,7 +19,6 @@
 
 import hopsworks_common.connection
 from hsfs.client import exceptions
-from hsfs.core import arrow_flight_client
 from hsfs.engine import spark, spark_no_metastore
 
 
@@ -83,5 +82,7 @@ def get_type() -> str:
 
 def stop() -> None:
     global _engine
+    from hsfs.core import arrow_flight_client
+
     _engine = None
     arrow_flight_client.close()

python/hsfs/engine/python.py

@@ -67,7 +67,6 @@
 from hsfs import storage_connector as sc
 from hsfs.constructor import query
 from hsfs.core import (
-    arrow_flight_client,
     dataset_api,
     feature_group_api,
     feature_view_api,
@@ -83,11 +82,12 @@
 )
 from hsfs.core.constants import (
     HAS_AIOMYSQL,
-    HAS_ARROW,
     HAS_GREAT_EXPECTATIONS,
     HAS_PANDAS,
+    HAS_PYARROW,
     HAS_SQLALCHEMY,
 )
+from hsfs.core.type_systems import PYARROW_HOPSWORKS_DTYPE_MAPPING
 from hsfs.core.vector_db_client import VectorDbClient
 from hsfs.feature_group import ExternalFeatureGroup, FeatureGroup
 from hsfs.training_dataset import TrainingDataset
@@ -98,8 +98,6 @@
 if HAS_GREAT_EXPECTATIONS:
     import great_expectations
 
-if HAS_ARROW:
-    from hsfs.core.type_systems import PYARROW_HOPSWORKS_DTYPE_MAPPING
 if HAS_AIOMYSQL and HAS_SQLALCHEMY:
     from hsfs.core import util_sql
 
@@ -157,6 +155,8 @@ def sql(
     def is_flyingduck_query_supported(
         self, query: "query.Query", read_options: Optional[Dict[str, Any]] = None
     ) -> bool:
+        from hsfs.core import arrow_flight_client
+
         return arrow_flight_client.is_query_supported(query, read_options or {})
 
     def _validate_dataframe_type(self, dataframe_type: str):
@@ -180,6 +180,8 @@ def _sql_offline(
     ) -> Union[pd.DataFrame, pl.DataFrame]:
         self._validate_dataframe_type(dataframe_type)
         if isinstance(sql_query, dict) and "query_string" in sql_query:
+            from hsfs.core import arrow_flight_client
+
             result_df = util.run_with_loading_animation(
                 "Reading data from Hopsworks, using Hopsworks Feature Query Service",
                 arrow_flight_client.get_instance().read_query,
@@ -342,6 +344,8 @@ def _read_hopsfs_remote(
 
             for inode in inode_list:
                 if not self._is_metadata_file(inode.path):
+                    from hsfs.core import arrow_flight_client
+
                     if arrow_flight_client.is_data_format_supported(
                         data_format, read_options
                     ):
@@ -539,7 +543,10 @@ def profile(
                 or pa.types.is_list(field.type)
                 or pa.types.is_large_list(field.type)
                 or pa.types.is_struct(field.type)
-            ) and PYARROW_HOPSWORKS_DTYPE_MAPPING[field.type] in ["timestamp", "date"]:
+            ) and PYARROW_HOPSWORKS_DTYPE_MAPPING.get(field.type, None) in [
+                "timestamp",
+                "date",
+            ]:
                 if HAS_POLARS and (
                     isinstance(df, pl.DataFrame)
                     or isinstance(df, pl.dataframe.frame.DataFrame)
@@ -573,15 +580,21 @@ def profile(
                 or pa.types.is_list(arrow_type)
                 or pa.types.is_large_list(arrow_type)
                 or pa.types.is_struct(arrow_type)
-                or PYARROW_HOPSWORKS_DTYPE_MAPPING[arrow_type]
+                or PYARROW_HOPSWORKS_DTYPE_MAPPING.get(arrow_type, None)
                 in ["timestamp", "date", "binary", "string"]
             ):
                 dataType = "String"
-            elif PYARROW_HOPSWORKS_DTYPE_MAPPING[arrow_type] in ["float", "double"]:
+            elif PYARROW_HOPSWORKS_DTYPE_MAPPING.get(arrow_type, None) in [
+                "float",
+                "double",
+            ]:
                 dataType = "Fractional"
-            elif PYARROW_HOPSWORKS_DTYPE_MAPPING[arrow_type] in ["int", "bigint"]:
+            elif PYARROW_HOPSWORKS_DTYPE_MAPPING.get(arrow_type, None) in [
+                "int",
+                "bigint",
+            ]:
                 dataType = "Integral"
-            elif PYARROW_HOPSWORKS_DTYPE_MAPPING[arrow_type] == "boolean":
+            elif PYARROW_HOPSWORKS_DTYPE_MAPPING.get(arrow_type, None) == "boolean":
                 dataType = "Boolean"
             else:
                 print(
@@ -1077,8 +1090,16 @@ def write_training_dataset(
                 "Currently only query based training datasets are supported by the Python engine"
             )
 
+        try:
+            from hsfs.core import arrow_flight_client
+
+            arrow_flight_client_imported = True
+        except ImportError:
+            arrow_flight_client_imported = False
+
         if (
-            arrow_flight_client.is_query_supported(dataset, user_write_options)
+            arrow_flight_client_imported
+            and arrow_flight_client.is_query_supported(dataset, user_write_options)
             and len(training_dataset.splits) == 0
             and feature_view_obj
             and len(feature_view_obj.transformation_functions) == 0
@@ -1251,7 +1272,7 @@ def _apply_transformation_function(
             or isinstance(dataset, pl.dataframe.frame.DataFrame)
         ):
             # Converting polars dataframe to pandas because currently we support only pandas UDF's as transformation functions.
-            if HAS_ARROW:
+            if HAS_PYARROW:
                 dataset = dataset.to_pandas(
                     use_pyarrow_extension_array=True
                 )  # Zero copy if pyarrow extension can be used.

python/hsfs/feature_store.py

@@ -37,7 +37,6 @@
 from hsfs.client import exceptions
 from hsfs.constructor.query import Query
 from hsfs.core import (
-    arrow_flight_client,
     feature_group_api,
     feature_group_engine,
     feature_view_engine,
@@ -1790,10 +1789,14 @@ def get_feature_views(self, name: str) -> List[feature_view.FeatureView]:
 
     def _disable_hopsworks_feature_query_service_client(self):
         """Disable Hopsworks feature query service for the current session. This behaviour is not persisted on reset."""
+        from hsfs.core import arrow_flight_client
+
         arrow_flight_client._disable_feature_query_service_client()
 
     def _reset_hopsworks_feature_query_service_client(self):
         """Reset Hopsworks feature query service for the current session."""
+        from hsfs.core import arrow_flight_client
+
         arrow_flight_client.close()
         arrow_flight_client.get_instance()
 

python/tests/core/test_type_systems.py

@@ -17,10 +17,10 @@
 
 import pytest
 from hsfs.core import type_systems
-from hsfs.core.constants import HAS_ARROW, HAS_PANDAS
+from hsfs.core.constants import HAS_PANDAS, HAS_PYARROW
 
 
-if HAS_ARROW:
+if HAS_PYARROW:
     import pyarrow as pa
 
 if HAS_PANDAS:
@@ -32,7 +32,7 @@
 
 class TestTypeSystems:
     @pytest.mark.skipif(
-        not HAS_ARROW or not HAS_PANDAS, reason="Arrow or Pandas are not installed"
+        not HAS_PYARROW or not HAS_PANDAS, reason="Arrow or Pandas are not installed"
     )
     def test_infer_type_pyarrow_list(self):
         # Act