Skip to content

Commit 3268ced

Browse files
makinzmmakinzm
authored
feat: Add versioning support to Milvus online store (#6330)
* feat: add versioning support to Milvus online store Signed-off-by: makinzm <nozomi.maki.da@gmail.com> * lint: uv run ruff format sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py Signed-off-by: makinzm <nozomi.maki.da@gmail.com> * feat: drop all versioned Milvus collections on teardown/update and lint Signed-off-by: makinzm <nozomi.maki.da@gmail.com> * chore: fix type to check whether it is None Signed-off-by: makinzm <nozomi.maki.da@gmail.com> --------- Signed-off-by: makinzm <nozomi.maki.da@gmail.com>
1 parent 5580ab4 commit 3268ced

3 files changed

Lines changed: 226 additions & 13 deletions

File tree

‎sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py‎

Lines changed: 38 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
deserialize_entity_key,
2020
serialize_entity_key,
2121
)
22+
from feast.infra.online_stores.helpers import compute_table_id
2223
from feast.infra.online_stores.online_store import OnlineStore
2324
from feast.infra.online_stores.vector_store import VectorStoreConfig
2425
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
@@ -164,7 +165,9 @@ def _get_or_create_collection(
164165
) -> Dict[str, Any]:
165166
self.client = self._connect(config)
166167
vector_field_dict = {k.name: k for k in table.schema if k.vector_index}
167-
collection_name = _table_id(config.project, table)
168+
collection_name = _table_id(
169+
config.project, table, config.registry.enable_online_feature_view_versioning
170+
)
168171
if collection_name not in self._collections:
169172
# Create a composite key by combining entity fields
170173
composite_key_name = _get_composite_key_name(table)
@@ -346,7 +349,9 @@ def online_read(
346349
requested_features: Optional[List[str]] = None,
347350
) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
348351
self.client = self._connect(config)
349-
collection_name = _table_id(config.project, table)
352+
collection_name = _table_id(
353+
config.project, table, config.registry.enable_online_feature_view_versioning
354+
)
350355
collection = self._get_or_create_collection(config, table)
351356

352357
composite_key_name = _get_composite_key_name(table)
@@ -493,11 +498,12 @@ def update(
493498
for table in tables_to_keep:
494499
self._get_or_create_collection(config, table)
495500

501+
# Always drop the base collection plus any "_v{N}" siblings, regardless of
502+
# the current versioning flag. This handles mixed-state repos where
503+
# versioning was toggled on/off across applies and would otherwise leave
504+
# orphan collections behind in Milvus.
496505
for table in tables_to_delete:
497-
collection_name = _table_id(config.project, table)
498-
if self._collections.get(collection_name, None):
499-
self.client.drop_collection(collection_name)
500-
self._collections.pop(collection_name, None)
506+
self._drop_all_version_collections(config.project, table)
501507

502508
def plan(
503509
self, config: RepoConfig, desired_registry_proto: RegistryProto
@@ -511,11 +517,9 @@ def teardown(
511517
entities: Sequence[Entity],
512518
):
513519
self.client = self._connect(config)
520+
# See update(): drop base + all "_v{N}" siblings to handle mixed-state repos.
514521
for table in tables:
515-
collection_name = _table_id(config.project, table)
516-
if self._collections.get(collection_name, None):
517-
self.client.drop_collection(collection_name)
518-
self._collections.pop(collection_name, None)
522+
self._drop_all_version_collections(config.project, table)
519523

520524
def retrieve_online_documents_v2(
521525
self,
@@ -551,7 +555,9 @@ def retrieve_online_documents_v2(
551555
k.name: k.dtype for k in table.entity_columns
552556
}
553557
self.client = self._connect(config)
554-
collection_name = _table_id(config.project, table)
558+
collection_name = _table_id(
559+
config.project, table, config.registry.enable_online_feature_view_versioning
560+
)
555561
collection = self._get_or_create_collection(config, table)
556562
if not config.online_store.vector_enabled:
557563
raise ValueError("Vector search is not enabled in the online store config")
@@ -748,9 +754,28 @@ def retrieve_online_documents_v2(
748754
result_list.append((res_ts, entity_key_proto, res if res else None))
749755
return result_list
750756

757+
def _drop_all_version_collections(self, project: str, table: FeatureView) -> None:
758+
"""Drop the base collection and every ``_v{N}`` versioned sibling.
759+
760+
Mirrors the ``_drop_all_version_tables`` helpers in the MySQL/PostgreSQL
761+
online stores. Always called from ``update`` and ``teardown`` so a
762+
repo that toggles versioning on and off does not leave orphan
763+
collections behind in Milvus.
764+
"""
765+
base = f"{project}_{table.name}"
766+
versioned_prefix = f"{base}_v"
767+
assert self.client is not None, "Milvus client is not initialized"
768+
for collection_name in self.client.list_collections():
769+
if collection_name == base or (
770+
collection_name.startswith(versioned_prefix)
771+
and collection_name[len(versioned_prefix) :].isdigit()
772+
):
773+
self.client.drop_collection(collection_name)
774+
self._collections.pop(collection_name, None)
775+
751776

752-
def _table_id(project: str, table: FeatureView) -> str:
753-
return f"{project}_{table.name}"
777+
def _table_id(project: str, table: FeatureView, enable_versioning: bool = False) -> str:
778+
return compute_table_id(project, table, enable_versioning)
754779

755780

756781
def _get_composite_key_name(table: FeatureView) -> str:

‎sdk/python/feast/infra/online_stores/online_store.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -292,6 +292,14 @@ def _check_versioned_read_support(self, grouped_refs):
292292
supported_types.append(DynamoDBOnlineStore)
293293
except Exception:
294294
pass
295+
try:
296+
from feast.infra.online_stores.milvus_online_store.milvus import (
297+
MilvusOnlineStore,
298+
)
299+
300+
supported_types.append(MilvusOnlineStore)
301+
except ImportError:
302+
pass
295303

296304
if isinstance(self, tuple(supported_types)):
297305
return

‎sdk/python/tests/unit/infra/online_store/test_milvus_versioning.py‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
"""Unit tests for Milvus online store feature view versioning."""
2+
3+
from datetime import timedelta
4+
from unittest.mock import MagicMock
5+
6+
from feast import Entity, FeatureView
7+
from feast.field import Field
8+
from feast.types import Float32
9+
from feast.value_type import ValueType
10+
11+
12+
def _make_feature_view(name="driver_stats", version_number=None, version_tag=None):
13+
entity = Entity(
14+
name="driver_id",
15+
join_keys=["driver_id"],
16+
value_type=ValueType.INT64,
17+
)
18+
fv = FeatureView(
19+
name=name,
20+
entities=[entity],
21+
ttl=timedelta(days=1),
22+
schema=[Field(name="trips_today", dtype=Float32)],
23+
)
24+
if version_number is not None:
25+
fv.current_version_number = version_number
26+
if version_tag is not None:
27+
fv.projection.version_tag = version_tag
28+
return fv
29+
30+
31+
def _make_config(project="test_project", versioning=False):
32+
config = MagicMock()
33+
config.project = project
34+
config.entity_key_serialization_version = 2
35+
config.registry.enable_online_feature_view_versioning = versioning
36+
return config
37+
38+
39+
class TestTableId:
40+
"""Test _table_id with versioning enabled/disabled."""
41+
42+
def test_no_versioning(self):
43+
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
44+
45+
fv = _make_feature_view()
46+
config = _make_config(versioning=False)
47+
assert _table_id(config.project, fv) == "test_project_driver_stats"
48+
49+
def test_versioning_enabled_with_version(self):
50+
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
51+
52+
fv = _make_feature_view(version_number=2)
53+
config = _make_config(versioning=True)
54+
assert (
55+
_table_id(config.project, fv, enable_versioning=True)
56+
== "test_project_driver_stats_v2"
57+
)
58+
59+
def test_projection_version_tag_takes_priority(self):
60+
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
61+
62+
fv = _make_feature_view(version_number=1, version_tag=3)
63+
config = _make_config(versioning=True)
64+
assert (
65+
_table_id(config.project, fv, enable_versioning=True)
66+
== "test_project_driver_stats_v3"
67+
)
68+
69+
def test_version_zero_no_suffix(self):
70+
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
71+
72+
fv = _make_feature_view(version_number=0)
73+
config = _make_config(versioning=True)
74+
assert (
75+
_table_id(config.project, fv, enable_versioning=True)
76+
== "test_project_driver_stats"
77+
)
78+
79+
def test_versioning_enabled_no_version_set(self):
80+
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
81+
82+
fv = _make_feature_view()
83+
config = _make_config(versioning=True)
84+
assert (
85+
_table_id(config.project, fv, enable_versioning=True)
86+
== "test_project_driver_stats"
87+
)
88+
89+
def test_versioning_disabled_ignores_version(self):
90+
from feast.infra.online_stores.milvus_online_store.milvus import _table_id
91+
92+
fv = _make_feature_view(version_number=5)
93+
config = _make_config(versioning=False)
94+
assert _table_id(config.project, fv) == "test_project_driver_stats"
95+
96+
97+
class TestMilvusVersionedReadSupport:
98+
"""Test that MilvusOnlineStore passes _check_versioned_read_support."""
99+
100+
def test_allowed_with_version_tag(self):
101+
from feast.infra.online_stores.milvus_online_store.milvus import (
102+
MilvusOnlineStore,
103+
)
104+
105+
store = MilvusOnlineStore()
106+
fv = _make_feature_view()
107+
fv.projection.version_tag = 2
108+
store._check_versioned_read_support([(fv, ["trips_today"])])
109+
110+
def test_allowed_without_version_tag(self):
111+
from feast.infra.online_stores.milvus_online_store.milvus import (
112+
MilvusOnlineStore,
113+
)
114+
115+
store = MilvusOnlineStore()
116+
fv = _make_feature_view()
117+
store._check_versioned_read_support([(fv, ["trips_today"])])
118+
119+
120+
class TestTeardownDropsAllVersions:
121+
"""Teardown should drop the base collection AND all versioned collections."""
122+
123+
def _build_store_with_collections(self, existing_collections):
124+
from feast.infra.online_stores.milvus_online_store.milvus import (
125+
MilvusOnlineStore,
126+
)
127+
128+
store = MilvusOnlineStore()
129+
store.client = MagicMock()
130+
store.client.list_collections.return_value = existing_collections
131+
store._connect = MagicMock(return_value=store.client)
132+
store._collections = {name: MagicMock() for name in existing_collections}
133+
return store
134+
135+
def test_teardown_drops_base_and_all_versioned_collections(self):
136+
fv = _make_feature_view()
137+
config = _make_config(versioning=True)
138+
existing = [
139+
"test_project_driver_stats",
140+
"test_project_driver_stats_v1",
141+
"test_project_driver_stats_v2",
142+
"test_project_other_view", # unrelated, must not be dropped
143+
]
144+
store = self._build_store_with_collections(existing)
145+
146+
store.teardown(config, [fv], [])
147+
148+
dropped = {call.args[0] for call in store.client.drop_collection.call_args_list}
149+
assert dropped == {
150+
"test_project_driver_stats",
151+
"test_project_driver_stats_v1",
152+
"test_project_driver_stats_v2",
153+
}
154+
assert "test_project_other_view" not in dropped
155+
156+
def test_update_drops_all_versions_for_deleted_table(self):
157+
fv = _make_feature_view()
158+
config = _make_config(versioning=True)
159+
existing = [
160+
"test_project_driver_stats",
161+
"test_project_driver_stats_v3",
162+
"test_project_driver_stats_v4",
163+
]
164+
store = self._build_store_with_collections(existing)
165+
166+
store.update(
167+
config=config,
168+
tables_to_delete=[fv],
169+
tables_to_keep=[],
170+
entities_to_delete=[],
171+
entities_to_keep=[],
172+
partial=False,
173+
)
174+
175+
dropped = {call.args[0] for call in store.client.drop_collection.call_args_list}
176+
assert dropped == {
177+
"test_project_driver_stats",
178+
"test_project_driver_stats_v3",
179+
"test_project_driver_stats_v4",
180+
}

0 commit comments

Comments
 (0)