Skip to content

Commit f9e91fa

Browse files
simonhearneclaude
simonhearne
and
claude
authored
fix: Load Milvus collections once instead of on every query (#6882)
* fix: Make Milvus placeholder vectors valid on Milvus servers Feature views without a vector field get a placeholder vector because Milvus requires one per collection. The placeholder was 1-dimensional, filled with NaN and left unindexed. Milvus Lite accepts all three, but Milvus servers and Zilliz Cloud reject dim < 2 and non-finite values, and refuse to load a collection with an unindexed vector field, so these feature views could not be created or read. The placeholder is now 2-dimensional, zero-filled and FLAT-indexed. Vector fields without vector_index=True also get a FLAT index so the collection can be loaded. Existing collections keep working because the placeholder is filled to the dimension recorded in the collection. Also fixes the Milvus docs page title and example config. Co-Authored-By: Claude Opus 5.5 <[email protected]> Signed-off-by: Simon Hearne <[email protected]> * fix: Load Milvus collections once instead of on every query online_read and retrieve_online_documents_v2 called load_collection on every request, adding a round-trip to each read. The call also hid the fact that newly created collections were never loaded on Milvus servers, because indexes were created after the collection. Collections are now created with their index params, which makes Milvus load them straight away. Existing collections are loaded only if their load state isn't Loaded, and only the first time a store accesses them. The per-query load_collection calls are removed. Co-Authored-By: Claude Opus 5.5 <[email protected]> Signed-off-by: Simon Hearne <[email protected]> --------- Signed-off-by: Simon Hearne <[email protected]> Co-authored-by: Claude Opus 5.5 <[email protected]>
1 parent e2d7fe2 commit f9e91fa

5 files changed

Lines changed: 691 additions & 34 deletions

File tree

‎.secrets.baseline‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

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

Lines changed: 48 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Redis online store
1+
# Milvus online store
22

33
## Description
44

@@ -19,11 +19,11 @@ Feast supports both milvus-lite 2.x and 3.x. However, if you upgrade from milvus
1919
See the [milvus-lite GitHub page](https://github.com/milvus-io/milvus-lite) for more details.
2020
{% endhint %}
2121

22-
You can get started by using any of the other templates (e.g. `feast init -t gcp` or `feast init -t snowflake` or `feast init -t aws`), and then swapping in Redis as the online store as seen below in the examples.
22+
You can get started by using any of the other templates (e.g. `feast init -t gcp` or `feast init -t snowflake` or `feast init -t aws`), and then swapping in Milvus as the online store as seen below in the examples.
2323

2424
## Examples
2525

26-
Connecting to a local MilvusDB instance:
26+
Using Milvus Lite, which stores data in a local file:
2727

2828
{% code title="feature_store.yaml" %}
2929
```yaml
@@ -33,18 +33,62 @@ provider: local
3333
online_store:
3434
type: milvus
3535
path: "data/online_store.db"
36-
connection_string: "localhost:6379"
3736
embedding_dim: 128
3837
index_type: "FLAT"
3938
metric_type: "COSINE"
39+
```
40+
{% endcode %}
41+
42+
Connecting to a self-hosted Milvus server:
43+
44+
{% code title="feature_store.yaml" %}
45+
```yaml
46+
project: my_feature_repo
47+
registry: data/registry.db
48+
provider: local
49+
online_store:
50+
type: milvus
51+
host: "http://localhost"
52+
port: 19530
4053
username: "username"
4154
password: "password"
55+
embedding_dim: 128
56+
index_type: "IVF_FLAT"
57+
metric_type: "COSINE"
4258
```
4359
{% endcode %}
4460
61+
## Configuration options
62+
63+
| Option | Default | Description |
64+
|:-------|:--------|:------------|
65+
| `path` | `""` | Path to a Milvus Lite database file. Used when `provider: local` and `path` is set. |
66+
| `host` | `http://localhost` | Milvus server host, including the scheme. |
67+
| `port` | `19530` | Milvus server port. |
68+
| `username` / `password` | `""` | Credentials, sent as the token `username:password`. |
69+
| `embedding_dim` | `128` | Dimension of vector fields. |
70+
| `index_type` | `FLAT` | Index type for vector fields with `vector_index=True`. |
71+
| `metric_type` | `COSINE` | Default metric when a field does not set `vector_search_metric`. |
72+
| `nlist` | `128` | `nlist` index parameter. |
73+
| `vector_enabled` | `true` | Enables vector search. |
74+
| `varchar_max_length` | `65535` | Default `max_length` of VARCHAR fields. Override per field with the `max_length` tag. |
75+
| `enable_openai_compatible_store` | `false` | Store numeric features as native Milvus numeric types. |
4576

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

79+
## Collection loading
80+
81+
Feast creates collections together with their indexes, which makes Milvus load them straight away.
82+
When Feast finds an existing collection it checks its load state and loads it only if needed.
83+
Reads and searches never load collections, so a collection released outside Feast is only reloaded
84+
the next time a Feast process first accesses it.
85+
86+
## Feature views without vectors
87+
88+
Milvus requires every collection to have a vector field. For feature views that have no vector
89+
feature, Feast adds a 2-dimensional `_placeholder_vector` field with a FLAT index and fills it with zeros.
90+
It is never returned or searched.
91+
4892
## Functionality Matrix
4993

5094
The set of functionality supported by online stores is described in detail [here](overview.md#functionality).

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

Lines changed: 59 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
FieldSchema,
1313
MilvusClient,
1414
)
15+
from pymilvus.client.types import LoadState
1516

1617
from feast import Entity
1718
from feast.feature_view import FeatureView
@@ -114,6 +115,13 @@
114115
DataType.BOOL,
115116
}
116117

118+
# Milvus requires every collection to have a vector field, so feature views
119+
# without one get a small placeholder vector. Milvus servers reject vectors
120+
# with fewer than 2 dimensions and non-finite values, and a collection can only
121+
# be loaded once every vector field is indexed.
122+
PLACEHOLDER_VECTOR_FIELD = "_placeholder_vector"
123+
PLACEHOLDER_VECTOR_DIM = 2
124+
117125

118126
def _milvus_escape_string(s: str) -> str:
119127
"""Escape a string for safe use inside a Milvus single-quoted literal.
@@ -348,9 +356,9 @@ def _get_or_create_collection(
348356
if not has_vector_field:
349357
fields.append(
350358
FieldSchema(
351-
name="_placeholder_vector",
359+
name=PLACEHOLDER_VECTOR_FIELD,
352360
dtype=DataType.FLOAT_VECTOR,
353-
dim=1,
361+
dim=PLACEHOLDER_VECTOR_DIM,
354362
)
355363
)
356364
schema = CollectionSchema(
@@ -360,22 +368,14 @@ def _get_or_create_collection(
360368
collection_name=collection_name
361369
)
362370
if not collection_exists:
363-
self.client.create_collection(
364-
collection_name=collection_name,
365-
dimension=config.online_store.embedding_dim,
366-
schema=schema,
367-
)
368371
index_params = self.client.prepare_index_params()
369-
indices_added = False
370372
for vector_field in schema.fields:
371-
if (
372-
vector_field.dtype
373-
in [
374-
DataType.FLOAT_VECTOR,
375-
DataType.BINARY_VECTOR,
376-
]
377-
and vector_field.name in vector_field_dict
378-
):
373+
if vector_field.dtype not in [
374+
DataType.FLOAT_VECTOR,
375+
DataType.BINARY_VECTOR,
376+
]:
377+
continue
378+
if vector_field.name in vector_field_dict:
379379
metric = vector_field_dict[
380380
vector_field.name
381381
].vector_search_metric
@@ -387,19 +387,44 @@ def _get_or_create_collection(
387387
index_name=f"vector_index_{vector_field.name}",
388388
params={"nlist": config.online_store.nlist},
389389
)
390-
indices_added = True
391-
if indices_added:
392-
self.client.create_index(
393-
collection_name=collection_name,
394-
index_params=index_params,
395-
)
390+
else:
391+
# Vector fields that aren't searched (the placeholder,
392+
# or arrays without vector_index) still need an index,
393+
# otherwise Milvus servers refuse to load the collection.
394+
index_params.add_index(
395+
collection_name=collection_name,
396+
field_name=vector_field.name,
397+
metric_type="L2"
398+
if vector_field.name == PLACEHOLDER_VECTOR_FIELD
399+
else config.online_store.metric_type,
400+
index_type="FLAT",
401+
index_name=f"vector_index_{vector_field.name}",
402+
)
403+
# Every collection has at least one vector field, and every
404+
# vector field is indexed, so passing the index params here
405+
# makes Milvus create the indexes and load the collection.
406+
self.client.create_collection(
407+
collection_name=collection_name,
408+
dimension=config.online_store.embedding_dim,
409+
schema=schema,
410+
index_params=index_params,
411+
)
396412
else:
397-
self.client.load_collection(collection_name)
413+
self._ensure_loaded(collection_name)
414+
# Collections are only cached once loaded, so reads and searches
415+
# don't need to load them again.
398416
self._collections[collection_name] = self.client.describe_collection(
399417
collection_name
400418
)
401419
return self._collections[collection_name]
402420

421+
def _ensure_loaded(self, collection_name: str) -> None:
422+
"""Load an existing collection unless Milvus already has it loaded."""
423+
assert self.client is not None, "Milvus client is not initialized"
424+
load_state = self.client.get_load_state(collection_name).get("state")
425+
if load_state != LoadState.Loaded:
426+
self.client.load_collection(collection_name)
427+
403428
def online_write_batch(
404429
self,
405430
config: RepoConfig,
@@ -423,6 +448,15 @@ def online_write_batch(
423448
collection_field_types = {
424449
field["name"]: field["type"] for field in collection["fields"]
425450
}
451+
# Collections created by older Feast versions may use a 1-dim placeholder.
452+
placeholder_dim = next(
453+
(
454+
int(field.get("params", {}).get("dim", PLACEHOLDER_VECTOR_DIM))
455+
for field in collection["fields"]
456+
if field["name"] == PLACEHOLDER_VECTOR_FIELD
457+
),
458+
PLACEHOLDER_VECTOR_DIM,
459+
)
426460
schema_internal_fields = {"event_ts", "created_ts"}
427461
collection_has_native_numerics = any(
428462
collection_field_types.get(name) in MILVUS_NATIVE_NUMERIC_TYPES
@@ -472,8 +506,8 @@ def online_write_batch(
472506
for field in required_fields:
473507
if field not in single_entity_record:
474508
field_type = collection_field_types.get(field, DataType.VARCHAR)
475-
if field == "_placeholder_vector":
476-
single_entity_record[field] = [float("nan")]
509+
if field == PLACEHOLDER_VECTOR_FIELD:
510+
single_entity_record[field] = [0.0] * placeholder_dim
477511
else:
478512
single_entity_record[field] = _default_for_milvus_type(
479513
field_type
@@ -533,7 +567,6 @@ def online_read(
533567
+ ", ".join([f"'{e}'" for e in composite_entities])
534568
+ "]"
535569
)
536-
self.client.load_collection(collection_name)
537570
results = self.client.query(
538571
collection_name=collection_name,
539572
filter=query_filter_for_entities,
@@ -746,8 +779,6 @@ def retrieve_online_documents_v2(
746779
ann_search_field = field["name"]
747780
break
748781

749-
self.client.load_collection(collection_name)
750-
751782
if filters and filters_contain_numeric_comparison(filters):
752783
collection_field_types = {
753784
f["name"]: f["type"] for f in collection["fields"]

0 commit comments

Comments
 (0)