Skip to content

Commit ef743ff

Browse files
simonhearneclaude
authored and
ntkathole
committed
feat: Make Milvus index and search params configurable
Search params were hard-coded to {"nprobe": 10} and index params to {"nlist": nlist}. index_type: AUTOINDEX failed on Milvus servers and Zilliz Cloud with "only metric type can be passed when use AutoIndex", because nlist and the collection name were sent as index params. Adds optional index_params and search_params config fields that are passed through to Milvus. When unset, the previous defaults are kept, except that AUTOINDEX gets no params. The collection name is no longer passed as an index param. AUTOINDEX can now be tuned with the level search param. Co-Authored-By: Claude Opus 5.5 <[email protected]> Signed-off-by: Simon Hearne <[email protected]>
1 parent 6161766 commit ef743ff

4 files changed

Lines changed: 208 additions & 7 deletions

File tree

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

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,13 +91,45 @@ online_store:
9191
| `embedding_dim` | `128` | Dimension of vector fields. |
9292
| `index_type` | `FLAT` | Index type for vector fields with `vector_index=True`. |
9393
| `metric_type` | `COSINE` | Default metric when a field does not set `vector_search_metric`. |
94-
| `nlist` | `128` | `nlist` index parameter. |
94+
| `nlist` | `128` | `nlist` index parameter, used when `index_params` is unset. |
95+
| `index_params` | unset | Index build parameters passed to Milvus, e.g. `{M: 16, efConstruction: 200}` for HNSW. Defaults to `{nlist: <nlist>}`, or no parameters for `AUTOINDEX`. |
96+
| `search_params` | unset | Search parameters passed to Milvus, e.g. `{ef: 64}` for HNSW or `{level: 2}` for `AUTOINDEX`. Defaults to `{nprobe: 10}`, or no parameters for `AUTOINDEX`. |
9597
| `vector_enabled` | `true` | Enables vector search. |
9698
| `varchar_max_length` | `65535` | Default `max_length` of VARCHAR fields. Override per field with the `max_length` tag. |
9799
| `enable_openai_compatible_store` | `false` | Store numeric features as native Milvus numeric types. |
98100

99101
The full set of configuration options is available in [MilvusOnlineStoreConfig](https://rtd.feast.dev/en/latest/#feast.infra.online_stores.milvus.MilvusOnlineStoreConfig).
100102

103+
## Index and search parameters
104+
105+
`index_type`, `index_params` and `search_params` are passed through to Milvus, so any index type the
106+
server supports can be used. On Zilliz Cloud, `AUTOINDEX` is recommended; tune the recall/latency
107+
trade-off with the `level` search parameter:
108+
109+
```yaml
110+
online_store:
111+
type: milvus
112+
index_type: "AUTOINDEX"
113+
search_params:
114+
level: 2
115+
```
116+
117+
For HNSW:
118+
119+
```yaml
120+
online_store:
121+
type: milvus
122+
index_type: "HNSW"
123+
index_params:
124+
M: 16
125+
efConstruction: 200
126+
search_params:
127+
ef: 64
128+
```
129+
130+
Index parameters only apply when Feast creates a collection. To change them for an existing
131+
collection, run `feast teardown` and `feast apply`, then materialize again.
132+
101133
## Collection loading
102134

103135
Feast creates collections together with their indexes, which makes Milvus load them straight away.

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

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,12 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
224224
vector_enabled: Optional[bool] = True
225225
text_search_enabled: Optional[bool] = False
226226
nlist: Optional[int] = 128
227+
# Index build params for vector fields, e.g. {"M": 16, "efConstruction": 200}.
228+
# Defaults to {"nlist": nlist}, or {} for AUTOINDEX.
229+
index_params: Optional[Dict[str, Any]] = None
230+
# Search params, e.g. {"ef": 64}, or {"level": 2} for AUTOINDEX.
231+
# Defaults to {"nprobe": 10}, or {} for AUTOINDEX.
232+
search_params: Optional[Dict[str, Any]] = None
227233
username: Optional[StrictStr] = ""
228234
password: Optional[StrictStr] = ""
229235
enable_openai_compatible_store: Optional[bool] = False
@@ -389,19 +395,17 @@ def _get_or_create_collection(
389395
vector_field.name
390396
].vector_search_metric
391397
index_params.add_index(
392-
collection_name=collection_name,
393398
field_name=vector_field.name,
394399
metric_type=metric or config.online_store.metric_type,
395400
index_type=config.online_store.index_type,
396401
index_name=f"vector_index_{vector_field.name}",
397-
params={"nlist": config.online_store.nlist},
402+
params=_index_build_params(config.online_store),
398403
)
399404
else:
400405
# Vector fields that aren't searched (the placeholder,
401406
# or arrays without vector_index) still need an index,
402407
# otherwise Milvus servers refuse to load the collection.
403408
index_params.add_index(
404-
collection_name=collection_name,
405409
field_name=vector_field.name,
406410
metric_type="L2"
407411
if vector_field.name == PLACEHOLDER_VECTOR_FIELD
@@ -848,7 +852,7 @@ def _combine_exprs(*parts: Optional[str]) -> Optional[str]:
848852

849853
search_params = {
850854
"metric_type": distance_metric or config.online_store.metric_type,
851-
"params": {"nprobe": 10},
855+
"params": _search_params(config.online_store),
852856
}
853857

854858
results = self.client.search(
@@ -865,7 +869,7 @@ def _combine_exprs(*parts: Optional[str]) -> Optional[str]:
865869
# Vector search only
866870
search_params = {
867871
"metric_type": distance_metric or config.online_store.metric_type,
868-
"params": {"nprobe": 10},
872+
"params": _search_params(config.online_store),
869873
}
870874

871875
results = self.client.search(
@@ -1028,6 +1032,27 @@ def _table_id(project: str, table: FeatureView, enable_versioning: bool = False)
10281032
return compute_table_id(project, table, enable_versioning)
10291033

10301034

1035+
def _is_autoindex(online_config: MilvusOnlineStoreConfig) -> bool:
1036+
return (online_config.index_type or "").upper() == "AUTOINDEX"
1037+
1038+
1039+
def _index_build_params(online_config: MilvusOnlineStoreConfig) -> Dict[str, Any]:
1040+
"""Build params for vector indexes. AUTOINDEX accepts none besides the metric."""
1041+
if online_config.index_params is not None:
1042+
return dict(online_config.index_params)
1043+
if _is_autoindex(online_config):
1044+
return {}
1045+
return {"nlist": online_config.nlist}
1046+
1047+
1048+
def _search_params(online_config: MilvusOnlineStoreConfig) -> Dict[str, Any]:
1049+
if online_config.search_params is not None:
1050+
return dict(online_config.search_params)
1051+
if _is_autoindex(online_config):
1052+
return {}
1053+
return {"nprobe": 10}
1054+
1055+
10311056
def _milvus_token(online_config: MilvusOnlineStoreConfig) -> str:
10321057
"""Return the token to authenticate with: ``token``, else ``username:password``."""
10331058
if online_config.token:

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

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -289,3 +289,26 @@ def test_db_name(tmp_path: Path, project: str, store: MilvusOnlineStore) -> None
289289
db_client = MilvusClient(uri=ZILLIZ_URI, token=ZILLIZ_TOKEN, db_name=db_name)
290290
assert collection_name in db_client.list_collections()
291291
assert collection_name not in admin.list_collections()
292+
293+
294+
def test_autoindex_with_search_level(
295+
tmp_path: Path, project: str, store: MilvusOnlineStore
296+
) -> None:
297+
config = _repo_config(
298+
tmp_path, project, index_type="AUTOINDEX", search_params={"level": 2}
299+
)
300+
fv = _vector_feature_view()
301+
store.update(config, [], [fv], [], [], partial=False)
302+
_write_rows(store, config, fv, _vector_rows())
303+
304+
assert store.client is not None
305+
index = store.client.describe_index(
306+
f"{project}_{fv.name}", "vector_index_embedding"
307+
)
308+
assert index["index_type"] == "AUTOINDEX"
309+
310+
hits = _eventually(
311+
lambda: _search(store, config, fv, [1.0, 0.0]),
312+
lambda hits: len(hits) == 1,
313+
)
314+
assert hits[0]["city"].string_val == "Paris"

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

Lines changed: 122 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
2121
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
2222
from feast.repo_config import RepoConfig
23-
from feast.types import Float32, Int64, String
23+
from feast.types import Array, Float32, Int64, String
2424
from feast.value_type import ValueType
2525

2626
MILVUS_MODULE = "feast.infra.online_stores.milvus_online_store.milvus"
@@ -375,3 +375,124 @@ def test_lite_path_used_without_uri(mock_client_cls: MagicMock) -> None:
375375
MilvusOnlineStore()._connect(config)
376376

377377
mock_client_cls.assert_called_once_with("/tmp/online_store.db")
378+
379+
380+
def _vector_feature_view(name: str = "driver_embeddings") -> FeatureView:
381+
return FeatureView(
382+
name=name,
383+
entities=[_driver_entity()],
384+
ttl=timedelta(days=1),
385+
schema=[
386+
Field(name="driver_id", dtype=Int64),
387+
Field(
388+
name="embedding",
389+
dtype=Array(Float32),
390+
vector_index=True,
391+
vector_search_metric="COSINE",
392+
),
393+
Field(name="city", dtype=String),
394+
],
395+
)
396+
397+
398+
def _vector_rows() -> Dict[int, Dict[str, ValueProto]]:
399+
def embedding(x: float, y: float) -> ValueProto:
400+
value = ValueProto()
401+
value.float_list_val.val.extend([x, y])
402+
return value
403+
404+
return {
405+
1: {"embedding": embedding(1.0, 0.0), "city": ValueProto(string_val="Paris")},
406+
2: {"embedding": embedding(0.0, 1.0), "city": ValueProto(string_val="Rome")},
407+
}
408+
409+
410+
def _search_with_mock(**online_store: Any) -> Dict[str, Any]:
411+
"""Create a collection and search it with a mocked client.
412+
413+
Returns the index definition of the embedding field and the search kwargs.
414+
"""
415+
with patch(f"{MILVUS_MODULE}.MilvusClient") as mock_client_cls:
416+
mock_client = _mock_client(mock_client_cls, has_collection=False)
417+
mock_client.describe_collection.return_value = {
418+
"collection_name": "test_milvus_driver_embeddings",
419+
"fields": [
420+
{"name": "driver_id_pk", "type": DataType.VARCHAR, "params": {}},
421+
{"name": "event_ts", "type": DataType.INT64, "params": {}},
422+
{"name": "created_ts", "type": DataType.INT64, "params": {}},
423+
{
424+
"name": "embedding",
425+
"type": DataType.FLOAT_VECTOR,
426+
"params": {"dim": 2},
427+
},
428+
{"name": "city", "type": DataType.VARCHAR, "params": {}},
429+
],
430+
}
431+
mock_client.search.return_value = [[]]
432+
433+
store = MilvusOnlineStore()
434+
store.retrieve_online_documents_v2(
435+
_mock_config(embedding_dim=2, **online_store),
436+
_vector_feature_view(),
437+
["embedding", "city"],
438+
embedding=[1.0, 0.0],
439+
top_k=1,
440+
)
441+
return {
442+
"index": _created_indexes(mock_client)["embedding"],
443+
"search": mock_client.search.call_args.kwargs,
444+
}
445+
446+
447+
def test_default_index_and_search_params_unchanged() -> None:
448+
result = _search_with_mock(index_type="IVF_FLAT")
449+
450+
assert result["index"]["index_type"] == "IVF_FLAT"
451+
assert result["index"]["nlist"] == 128
452+
assert result["search"]["search_params"]["params"] == {"nprobe": 10}
453+
454+
455+
def test_autoindex_gets_no_index_or_search_params_by_default() -> None:
456+
result = _search_with_mock(index_type="AUTOINDEX")
457+
458+
# Milvus servers reject AUTOINDEX with any build param besides the metric.
459+
assert result["index"] == {
460+
"field_name": "embedding",
461+
"index_type": "AUTOINDEX",
462+
"index_name": "vector_index_embedding",
463+
"metric_type": "COSINE",
464+
}
465+
assert result["search"]["search_params"]["params"] == {}
466+
467+
468+
def test_index_and_search_params_pass_through() -> None:
469+
result = _search_with_mock(
470+
index_type="HNSW",
471+
index_params={"M": 16, "efConstruction": 200},
472+
search_params={"ef": 64},
473+
)
474+
475+
assert result["index"]["M"] == 16
476+
assert result["index"]["efConstruction"] == 200
477+
assert "nlist" not in result["index"]
478+
assert result["search"]["search_params"]["params"] == {"ef": 64}
479+
480+
481+
def test_autoindex_search_with_level(tmp_path: Path) -> None:
482+
config = _lite_config(tmp_path, index_type="AUTOINDEX", search_params={"level": 2})
483+
fv = _vector_feature_view()
484+
store = MilvusOnlineStore()
485+
store.update(config, [], [fv], [], [], partial=False)
486+
_write_rows(store, config, fv, _vector_rows())
487+
488+
results = store.retrieve_online_documents_v2(
489+
config,
490+
fv,
491+
["embedding", "city"],
492+
embedding=[1.0, 0.0],
493+
top_k=1,
494+
distance_metric="COSINE",
495+
)
496+
497+
assert len(results) == 1
498+
assert results[0][2] is not None and results[0][2]["city"].string_val == "Paris"

0 commit comments

Comments
 (0)