ExpediaGroup · michaelbackes · Aug 15, 2023 · Aug 14, 2023 · Aug 15, 2023 · Aug 15, 2023
diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml
@@ -10,6 +10,8 @@ jobs:
         python-version: ["3.8", "3.9", "3.10"]
         os: [ubuntu-latest, macOS-latest]
         exclude:
+          - os: macOS-latest
+            python-version: "3.8"
           - os: macOS-latest
             python-version: "3.9"
           - os: macOS-latest

diff --git a/milvus/Dockerfile b/milvus/Dockerfile
@@ -0,0 +1,12 @@
+FROM python:3.11-slim
+
+RUN python3 -m pip install milvus==2.2.12
+
+# this is needed to divert logs to stdout
+RUN mkdir -p /root/.milvus.io/milvus-server/2.2.12/logs/
+RUN touch /root/.milvus.io/milvus-server/2.2.12/logs/milvus-stdout.log
+RUN touch /root/.milvus.io/milvus-server/2.2.12/logs/milvus-stderr.log
+RUN ln -sf /dev/stdout /root/.milvus.io/milvus-server/2.2.12/logs/milvus-stdout.log \
+    && ln -sf /dev/stderr /root/.milvus.io/milvus-server/2.2.12/logs/milvus-stderr.log
+
+CMD ["milvus-server"]
diff --git a/sdk/python/feast/diff/registry_diff.py b/sdk/python/feast/diff/registry_diff.py
@@ -133,7 +133,7 @@ def diff_registry_objects(
     if isinstance(
         current_proto, (DataSourceProto, ValidationReferenceProto)
     ) or isinstance(new_proto, (DataSourceProto, ValidationReferenceProto)):
-        assert type(current_proto) == type(new_proto)
+        assert type(current_proto) is type(new_proto)
         current_spec = cast(DataSourceProto, current_proto)
         new_spec = cast(DataSourceProto, new_proto)
     else:

diff --git a/sdk/python/feast/driver_test_data.py b/sdk/python/feast/driver_test_data.py
@@ -187,6 +187,10 @@ def create_customer_daily_profile_df(customers, start_date, end_date) -> pd.Data
     df_all_customers["lifetime_trip_count"] = np.random.randint(
         0, 1000, size=rows
     ).astype(np.int32)
+    df_all_customers["profile_embedding"] = [
+        np.random.default_rng().uniform(-100, 200, 50).astype(np.float32)
+        for _ in range(rows)
+    ]
 
     # TODO: Remove created timestamp in order to test whether its really optional
     df_all_customers["created"] = pd.to_datetime(pd.Timestamp.now(tz=None).round("ms"))

diff --git a/sdk/python/feast/expediagroup/vectordb/milvus_online_store.py b/sdk/python/feast/expediagroup/vectordb/milvus_online_store.py
@@ -63,11 +63,13 @@ def __init__(self, online_config: RepoConfig):
     def __enter__(self):
         # Connecting to Milvus
         logger.info(
-            f"Connecting to Milvus with alias {self.online_config.alias} and host {self.online_config.host} and default port {self.online_config.port}."
+            f"Connecting to Milvus with alias {self.online_config.alias} and host {self.online_config.host} and port {self.online_config.port}."
         )
         connections.connect(
+            alias=self.online_config.alias,
             host=self.online_config.host,
-            username=self.online_config.username,
+            port=self.online_config.port,
+            user=self.online_config.username,
             password=self.online_config.password,
             use_secure=True,
         )
@@ -158,9 +160,7 @@ def teardown(
         tables: Sequence[VectorFeatureView],
         entities: Sequence[Entity],
     ):
-        raise NotImplementedError(
-            "to be implemented in https://jira.expedia.biz/browse/EAPC-7974"
-        )
+        pass
 
     def _convert_featureview_schema_to_milvus_readable(
         self, feast_schema: List[Field], vector_field, vector_field_dimensions

diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -1109,7 +1109,7 @@ def get_historical_features(
         set_usage_attribute("request_fv", bool(request_feature_views))
 
         # Check that the right request data is present in the entity_df
-        if type(entity_df) == pd.DataFrame:
+        if isinstance(entity_df, pd.DataFrame):
             if self.config.coerce_tz_aware:
                 entity_df = utils.make_df_tzaware(cast(pd.DataFrame, entity_df))
             for fv in request_feature_views:

diff --git a/sdk/python/feast/field.py b/sdk/python/feast/field.py
@@ -65,7 +65,7 @@ def dtype_is_feasttype_or_string_feasttype(cls, v):
         return v
 
     def __eq__(self, other):
-        if type(self) != type(other):
+        if type(self) is not type(other):
             return False
 
         if (

diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py
@@ -725,7 +725,7 @@ def _get_entity_df_event_timestamp_range(
             or entity_df_event_timestamp_range[1] is None
         ):
             raise EntitySQLEmptyResults(entity_df)
-        if type(entity_df_event_timestamp_range[0]) != datetime:
+        if not isinstance(entity_df_event_timestamp_range[0], datetime):
             raise EntityDFNotDateTime()
     elif isinstance(entity_df, pd.DataFrame):
         entity_df_event_timestamp = entity_df.loc[

diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py
@@ -203,8 +203,6 @@ def __init__(self, **data: Any):
                 self._offline_config = "redshift"
             elif data["provider"] == "azure":
                 self._offline_config = "mssql"
-            elif data["provider"] == "milvus":
-                self._online_config = "milvus"
 
         self._online_store = None
         if "online_store" in data:
@@ -218,8 +216,6 @@ def __init__(self, **data: Any):
                 self._online_config = "dynamodb"
             elif data["provider"] == "rockset":
                 self._online_config = "rockset"
-            elif data["provider"] == "milvus":
-                self._online_config = "milvus"
 
         self._batch_engine = None
         if "batch_engine" in data:

diff --git a/sdk/python/feast/type_map.py b/sdk/python/feast/type_map.py
@@ -745,7 +745,7 @@ def spark_to_feast_value_type(spark_type_as_str: str) -> ValueType:
         "array<timestamp>": ValueType.UNIX_TIMESTAMP_LIST,
     }
     # TODO: Find better way of doing this.
-    if type(spark_type_as_str) != str or spark_type_as_str not in type_map:
+    if not isinstance(spark_type_as_str, str) or spark_type_as_str not in type_map:
         return ValueType.NULL
     return type_map[spark_type_as_str.lower()]
 

diff --git a/sdk/python/tests/expediagroup/milvus_online_store_creator.py b/sdk/python/tests/expediagroup/milvus_online_store_creator.py
@@ -1,6 +1,7 @@
 from typing import Dict
 
-from milvus import default_server
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.waiting_utils import wait_for_logs
 
 from tests.integration.feature_repos.universal.online_store_creator import (
     OnlineStoreCreator,
@@ -10,15 +11,28 @@
 class MilvusOnlineStoreCreator(OnlineStoreCreator):
     def __init__(self, project_name: str, **kwargs):
         super().__init__(project_name)
-        self.host = "localhost"
-        self.port = 19530
-        self.server = default_server
-        self.server.wait_for_started = False
+        self.container = DockerContainer(
+            "mbackes/milvus-lite:2.2.12"
+        ).with_exposed_ports("19530")
 
     def create_online_store(self) -> Dict[str, str]:
-        self.server.start()
+        self.container.start()
+        log_string_to_wait_for = (
+            "Milvus Proxy successfully initialized and ready to serve!"
+        )
+        wait_for_logs(
+            container=self.container, predicate=log_string_to_wait_for, timeout=30
+        )
+        exposed_port = self.container.get_exposed_port("19530")
 
-        return {"type": "milvus", "host": self.host, "port": str(self.port)}
+        return {
+            "alias": "default",
+            "type": "milvus",
+            "host": "localhost",
+            "port": str(exposed_port),
+            "username": "user",
+            "password": "password",
+        }
 
     def teardown(self):
-        self.server.stop()
+        self.container.stop()