Merge pull request #11 from ExpediaGroup/feature/vector_online_store

adding online store for Milvus

sdk/python/feast/expediagroup/vectordb/milvus_online_store.py

@@ -0,0 +1,73 @@
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+
+from pydantic.typing import Literal
+
+from feast import Entity, RepoConfig
+from feast.expediagroup.vectordb.vector_feature_view import VectorFeatureView
+from feast.expediagroup.vectordb.vector_online_store import VectorOnlineStore
+from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
+from feast.protos.feast.types.Value_pb2 import Value as ValueProto
+from feast.repo_config import FeastConfigBaseModel
+
+
+class MilvusOnlineStoreConfig(FeastConfigBaseModel):
+    """Online store config for the Milvus online store"""
+
+    type: Literal["milvus"] = "milvus"
+    """Online store type selector"""
+
+    host: str
+    """ the host URL """
+
+    port: int = 19530
+    """ the port to connect to a Milvus instance. Should be the one used for GRPC (default: 19530) """
+
+
+class MilvusOnlineStore(VectorOnlineStore):
+    def online_write_batch(
+        self,
+        config: RepoConfig,
+        table: VectorFeatureView,
+        data: List[
+            Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
+        ],
+        progress: Optional[Callable[[int], Any]],
+    ) -> None:
+        raise NotImplementedError(
+            "to be implemented in https://jira.expedia.biz/browse/EAPC-7971"
+        )
+
+    def online_read(
+        self,
+        config: RepoConfig,
+        table: VectorFeatureView,
+        entity_keys: List[EntityKeyProto],
+        requested_features: Optional[List[str]] = None,
+    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
+        raise NotImplementedError(
+            "to be implemented in https://jira.expedia.biz/browse/EAPC-7972"
+        )
+
+    def update(
+        self,
+        config: RepoConfig,
+        tables_to_delete: Sequence[VectorFeatureView],
+        tables_to_keep: Sequence[VectorFeatureView],
+        entities_to_delete: Sequence[Entity],
+        entities_to_keep: Sequence[Entity],
+        partial: bool,
+    ):
+        raise NotImplementedError(
+            "to be implemented in https://jira.expedia.biz/browse/EAPC-7970"
+        )
+
+    def teardown(
+        self,
+        config: RepoConfig,
+        tables: Sequence[VectorFeatureView],
+        entities: Sequence[Entity],
+    ):
+        raise NotImplementedError(
+            "to be implemented in https://jira.expedia.biz/browse/EAPC-7974"
+        )

sdk/python/feast/expediagroup/vectordb/vector_online_store.py

@@ -0,0 +1,63 @@
+from abc import abstractmethod
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
+
+from feast import Entity, RepoConfig
+from feast.expediagroup.vectordb.vector_feature_view import VectorFeatureView
+from feast.infra.online_stores.online_store import OnlineStore
+from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
+from feast.protos.feast.types.Value_pb2 import Value as ValueProto
+
+
+class VectorOnlineStore(OnlineStore):
+    """
+    Abstraction for vector database implementations of the online store interface. Any online store implementation of
+    a vector database should inherit this class.
+    """
+
+    @abstractmethod
+    def online_write_batch(
+        self,
+        config: RepoConfig,
+        table: VectorFeatureView,
+        data: List[
+            Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
+        ],
+        progress: Optional[Callable[[int], Any]],
+    ) -> None:
+        # to be implemented in inheriting class
+        pass
+
+    @abstractmethod
+    def online_read(
+        self,
+        config: RepoConfig,
+        table: VectorFeatureView,
+        entity_keys: List[EntityKeyProto],
+        requested_features: Optional[List[str]] = None,
+    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
+        # to be implemented in inheriting class
+        pass
+
+    @abstractmethod
+    def update(
+        self,
+        config: RepoConfig,
+        tables_to_delete: Sequence[VectorFeatureView],
+        tables_to_keep: Sequence[VectorFeatureView],
+        entities_to_delete: Sequence[Entity],
+        entities_to_keep: Sequence[Entity],
+        partial: bool,
+    ):
+        # to be implemented in inheriting class
+        pass
+
+    @abstractmethod
+    def teardown(
+        self,
+        config: RepoConfig,
+        tables: Sequence[VectorFeatureView],
+        entities: Sequence[Entity],
+    ):
+        # to be implemented in inheriting class
+        pass

sdk/python/feast/repo_config.py

@@ -64,6 +64,7 @@
     "mysql": "feast.infra.online_stores.contrib.mysql_online_store.mysql.MySQLOnlineStore",
     "rockset": "feast.infra.online_stores.contrib.rockset_online_store.rockset.RocksetOnlineStore",
     "hazelcast": "feast.infra.online_stores.contrib.hazelcast_online_store.hazelcast_online_store.HazelcastOnlineStore",
+    "milvus": "feast.expediagroup.vectordb.milvus_online_store.MilvusOnlineStore",
 }
 
 OFFLINE_STORE_CLASS_FOR_TYPE = {