-
Notifications
You must be signed in to change notification settings - Fork 3
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Msudhir/add vector update functionality #14
Changes from 25 commits
d372a09
f6d3caf
d4f9158
4861af0
48e0971
315073f
870762a
478caec
1c01035
ef4ef32
0ad2d62
e4c0c9b
bef5791
928be7b
12f57a9
9527183
76270f6
c75a01f
0578b9b
8487678
5828891
4a29d33
e209770
ebe1e32
62692e0
0680c94
5c5490d
cdadb87
e1fd230
d0c4269
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,14 +1,19 @@ | ||
import logging | ||
from datetime import datetime | ||
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple | ||
|
||
from pydantic.typing import Literal | ||
from pymilvus import Collection, connections, utility | ||
|
||
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 | ||
from feast.usage import log_exceptions_and_usage | ||
|
||
logger = logging.getLogger(__name__) | ||
|
||
|
||
class MilvusOnlineStoreConfig(FeastConfigBaseModel): | ||
|
@@ -17,13 +22,47 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel): | |
type: Literal["milvus"] = "milvus" | ||
"""Online store type selector""" | ||
|
||
alias: str = "default" | ||
""" alias for milvus connection""" | ||
|
||
host: str | ||
""" the host URL """ | ||
|
||
username: str | ||
""" username to connect to Milvus """ | ||
|
||
password: str | ||
""" password to connect to Milvus """ | ||
|
||
port: int = 19530 | ||
""" the port to connect to a Milvus instance. Should be the one used for GRPC (default: 19530) """ | ||
|
||
|
||
class MilvusConnectionManager: | ||
def __init__(self, online_config: RepoConfig): | ||
self.online_config = online_config | ||
|
||
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}." | ||
) | ||
connections.connect( | ||
host=self.online_config.host, | ||
username=self.online_config.username, | ||
password=self.online_config.password, | ||
use_secure=True, | ||
) | ||
|
||
def __exit__(self, exc_type, exc_value, traceback): | ||
# Disconnecting from Milvus | ||
logger.info("Closing the connection to Milvus") | ||
connections.disconnect(self.online_config.alias) | ||
logger.info("Connection Closed") | ||
if exc_type is not None: | ||
logger.error(f"An exception of type {exc_type} occurred: {exc_value}") | ||
|
||
|
||
class MilvusOnlineStore(VectorOnlineStore): | ||
def online_write_batch( | ||
self, | ||
|
@@ -49,6 +88,7 @@ def online_read( | |
"to be implemented in https://jira.expedia.biz/browse/EAPC-7972" | ||
) | ||
|
||
@log_exceptions_and_usage(online_store="milvus") | ||
def update( | ||
self, | ||
config: RepoConfig, | ||
|
@@ -58,9 +98,34 @@ def update( | |
entities_to_keep: Sequence[Entity], | ||
partial: bool, | ||
): | ||
raise NotImplementedError( | ||
"to be implemented in https://jira.expedia.biz/browse/EAPC-7970" | ||
) | ||
with MilvusConnectionManager(config.online_store): | ||
for table_to_keep in tables_to_keep: | ||
collection_available = utility.has_collection(table_to_keep.name) | ||
try: | ||
if collection_available: | ||
logger.error(f"Collection {table_to_keep.name} already exists.") | ||
else: | ||
Collection(name=table_to_keep.name, schema=table_to_keep.schema) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. just catching this. FeatureView schemas are not the same as the schema that Milvus expects. We will need to write logic to translate that. |
||
logger.info( | ||
f"Collection {table_to_keep.name} has been created successfully." | ||
) | ||
except Exception as e: | ||
logger.error(f"Collection update failed due to {e}") | ||
|
||
for table_to_delete in tables_to_delete: | ||
collection_available = utility.has_collection(table_to_delete.name) | ||
try: | ||
if collection_available: | ||
utility.drop_collection(table_to_delete.name) | ||
logger.info( | ||
f"Collection {table_to_delete.name} has been deleted successfully." | ||
) | ||
else: | ||
return logger.error( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. remove return statement |
||
f"Collection {table_to_delete.name} does not exist or is already deleted." | ||
) | ||
except Exception as e: | ||
logger.error(f"Collection deletion failed due to {e}") | ||
|
||
def teardown( | ||
self, | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
please make this log statement on "info" level. If a collection exists it is not an error and is expected.