Skip to content

Commit 68de84e

Browse files
jiwidishuchu
andauthored
feat: Add num_partitions and partition_key_isolation to the Milvus online store (#7083)
* feat: Add num_partitions and partition_key_isolation to the Milvus online store A partition key (#6940) comes with two collection settings that Milvus only accepts when the collection is created: the number of partitions, and partition key isolation, which makes a search filtered on the key scan only that key's data. num_partitions and partition_key_isolation are passed to create_collection for feature views that have a partition key. They are ignored for other feature views. When unset nothing is sent, so Milvus keeps its defaults. Signed-off-by: jiwidi <[email protected]> * ci: Re-run the checks after a macOS runner failure Signed-off-by: jiwidi <[email protected]> --------- Signed-off-by: jiwidi <[email protected]> Co-authored-by: Shuchu Han <[email protected]>
1 parent a56766e commit 68de84e

4 files changed

Lines changed: 132 additions & 0 deletions

File tree

‎docs/reference/online-stores/milvus.md‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,8 @@ online_store:
9797
| `consistency_level` | unset | `Strong`, `Bounded`, `Session` or `Eventually`. Sent with every read and search. When unset, Milvus uses the collection's level. |
9898
| `collection_consistency_level` | unset | `Strong`, `Bounded`, `Session` or `Eventually`. Set when Feast creates a collection. When unset, Milvus uses its default (`Bounded`). |
9999
| `partition_key` | unset | Field to use as the Milvus partition key in feature views that contain it. See [Partition key](#partition-key). |
100+
| `num_partitions` | unset | Number of partitions of a new collection that has a partition key. When unset, Milvus uses its default. |
101+
| `partition_key_isolation` | unset | Enables partition key isolation on a new collection that has a partition key. When unset, Milvus uses its default. |
100102
| `vector_enabled` | `true` | Enables vector search. |
101103
| `varchar_max_length` | `65535` | Default `max_length` of VARCHAR fields. Override per field with the `max_length` tag. |
102104
| `enable_openai_compatible_store` | `false` | Store numeric features as native Milvus numeric types. |
@@ -168,6 +170,12 @@ store.retrieve_online_documents_v2(
168170
)
169171
```
170172

173+
`num_partitions` sets how many partitions a new collection with a partition key gets, and
174+
`partition_key_isolation` enables
175+
[partition key isolation](https://milvus.io/docs/use-partition-key.md#Use-Partition-Key-Isolation)
176+
for it, which Milvus supports with the HNSW index. Both settings are ignored for feature views
177+
without a partition key, and Milvus Lite ignores `num_partitions`.
178+
171179
The partition key field must be stored as `VARCHAR` or `INT64`. String fields are always `VARCHAR`,
172180
and integer fields are `VARCHAR` unless native numeric types are enabled, in which case `Int64` works
173181
and `Int32` does not.

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

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,11 +250,24 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
250250
# Field to use as the Milvus partition key in feature views that contain it.
251251
# A feature view's "milvus.partition_key" tag takes precedence.
252252
partition_key: Optional[StrictStr] = None
253+
# Number of partitions of a new collection that has a partition key. When
254+
# unset, Milvus uses its default.
255+
num_partitions: Optional[int] = None
256+
# Partition key isolation of a new collection that has a partition key.
257+
# When unset, Milvus uses its default.
258+
partition_key_isolation: Optional[bool] = None
253259
username: Optional[StrictStr] = ""
254260
password: Optional[StrictStr] = ""
255261
enable_openai_compatible_store: Optional[bool] = False
256262
varchar_max_length: Optional[int] = 65535
257263

264+
@field_validator("num_partitions")
265+
@classmethod
266+
def validate_num_partitions(cls, v):
267+
if v is not None and v < 1:
268+
raise ValueError(f"num_partitions must be at least 1, got {v}")
269+
return v
270+
258271
@field_validator("varchar_max_length")
259272
@classmethod
260273
def validate_varchar_max_length(cls, v):
@@ -445,6 +458,7 @@ def _get_or_create_collection(
445458
schema=schema,
446459
index_params=index_params,
447460
**_collection_consistency_kwargs(config.online_store),
461+
**_partition_kwargs(config.online_store, partition_key),
448462
)
449463
else:
450464
self._ensure_loaded(collection_name)
@@ -1195,6 +1209,22 @@ def _consistency_kwargs(online_config: MilvusOnlineStoreConfig) -> Dict[str, Any
11951209
return {}
11961210

11971211

1212+
def _partition_kwargs(
1213+
online_config: MilvusOnlineStoreConfig, partition_key: Optional[str]
1214+
) -> Dict[str, Any]:
1215+
"""Partition settings only apply to a collection that has a partition key."""
1216+
kwargs: Dict[str, Any] = {}
1217+
if not partition_key:
1218+
return kwargs
1219+
if online_config.num_partitions is not None:
1220+
kwargs["num_partitions"] = online_config.num_partitions
1221+
if online_config.partition_key_isolation is not None:
1222+
kwargs["properties"] = {
1223+
"partitionkey.isolation": str(online_config.partition_key_isolation).lower()
1224+
}
1225+
return kwargs
1226+
1227+
11981228
def _collection_consistency_kwargs(
11991229
online_config: MilvusOnlineStoreConfig,
12001230
) -> Dict[str, Any]:

‎sdk/python/tests/integration/online_store/test_milvus_remote.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -409,6 +409,63 @@ def test_partition_key_filtering(
409409
assert [hit["city"].string_val for hit in hits] == ["Rome"]
410410

411411

412+
def test_partition_options(
413+
tmp_path: Path, project: str, store: MilvusOnlineStore
414+
) -> None:
415+
config = _repo_config(
416+
tmp_path,
417+
project,
418+
partition_key="brand_id",
419+
num_partitions=4,
420+
partition_key_isolation=True,
421+
index_type="HNSW",
422+
consistency_level="Strong",
423+
)
424+
fv = FeatureView(
425+
name="products",
426+
entities=[
427+
Entity(
428+
name="driver_id", join_keys=["driver_id"], value_type=ValueType.INT64
429+
)
430+
],
431+
ttl=timedelta(days=1),
432+
schema=[
433+
Field(name="driver_id", dtype=Int64),
434+
Field(name="brand_id", dtype=String),
435+
Field(
436+
name="embedding",
437+
dtype=Array(Float32),
438+
vector_index=True,
439+
vector_search_metric="COSINE",
440+
),
441+
Field(name="city", dtype=String),
442+
],
443+
)
444+
store.update(config, [], [fv], [], [], partial=False)
445+
446+
assert store.client is not None
447+
description = store.client.describe_collection(f"{project}_{fv.name}")
448+
assert description["num_partitions"] == 4
449+
properties = description["properties"]
450+
if isinstance(properties, list):
451+
properties = {entry["key"]: entry["value"] for entry in properties}
452+
assert properties["partitionkey.isolation"] == "true"
453+
454+
rows = _vector_rows()
455+
rows[1]["brand_id"] = ValueProto(string_val="acme")
456+
rows[2]["brand_id"] = ValueProto(string_val="globex")
457+
_write_rows(store, config, fv, rows)
458+
hits = _search(
459+
store,
460+
config,
461+
fv,
462+
[1.0, 0.0],
463+
top_k=5,
464+
filters=ComparisonFilter(type="eq", key="brand_id", value="globex"),
465+
)
466+
assert [hit["city"].string_val for hit in hits] == ["Rome"]
467+
468+
412469
def test_online_delete(tmp_path: Path, project: str, store: MilvusOnlineStore) -> None:
413470
config = _repo_config(tmp_path, project, consistency_level="Strong")
414471
fv = _vector_feature_view()

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -748,6 +748,43 @@ def test_warns_when_existing_collection_lacks_partition_key(
748748
mock_client.create_collection.assert_not_called()
749749

750750

751+
@patch(f"{MILVUS_MODULE}.MilvusClient")
752+
def test_partition_options_sent_only_for_partitioned_collections(
753+
mock_client_cls: MagicMock,
754+
) -> None:
755+
mock_client = _mock_client(mock_client_cls, has_collection=False)
756+
config = _mock_config(
757+
partition_key="brand_id", num_partitions=32, partition_key_isolation=True
758+
)
759+
store = MilvusOnlineStore()
760+
761+
store._get_or_create_collection(config, _catalog_feature_view())
762+
kwargs = mock_client.create_collection.call_args.kwargs
763+
assert kwargs["num_partitions"] == 32
764+
assert kwargs["properties"] == {"partitionkey.isolation": "true"}
765+
766+
store._get_or_create_collection(config, _scalar_feature_view())
767+
kwargs = mock_client.create_collection.call_args.kwargs
768+
assert "num_partitions" not in kwargs and "properties" not in kwargs
769+
770+
771+
@patch(f"{MILVUS_MODULE}.MilvusClient")
772+
def test_partition_options_not_sent_when_unset(mock_client_cls: MagicMock) -> None:
773+
mock_client = _mock_client(mock_client_cls, has_collection=False)
774+
775+
MilvusOnlineStore()._get_or_create_collection(
776+
_mock_config(partition_key="brand_id"), _catalog_feature_view()
777+
)
778+
779+
kwargs = mock_client.create_collection.call_args.kwargs
780+
assert "num_partitions" not in kwargs and "properties" not in kwargs
781+
782+
783+
def test_invalid_num_partitions_rejected() -> None:
784+
with pytest.raises(ValidationError, match="num_partitions"):
785+
MilvusOnlineStoreConfig(num_partitions=0)
786+
787+
751788
def test_online_delete_by_entity_key_and_by_filter(tmp_path: Path) -> None:
752789
config = _lite_config(tmp_path)
753790
fv = _scalar_feature_view()

0 commit comments

Comments
 (0)