Skip to content

fix: Load Milvus collections once instead of on every query #6882

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.

Sign up for GitHub

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

Merged
ntkathole merged 2 commits into from
Sep 29, 2026
Merged

fix: Load Milvus collections once instead of on every query #6882

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0838f0d
fix: Make Milvus placeholder vectors valid on Milvus servers
simonhearne Sep 28, 2026
f75c444
fix: Load Milvus collections once instead of on every query
simonhearne Sep 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .secrets.baseline
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

52 changes: 48 additions & 4 deletions docs/reference/online-stores/milvus.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# Redis online store
# Milvus online store

## Description

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

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.
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.

## Examples

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

{% code title="feature_store.yaml" %}
```yaml
Expand All @@ -33,18 +33,62 @@ provider: local
online_store:
type: milvus
path: "data/online_store.db"
connection_string: "localhost:6379"
embedding_dim: 128
index_type: "FLAT"
metric_type: "COSINE"
```
{% endcode %}

Connecting to a self-hosted Milvus server:

{% code title="feature_store.yaml" %}
```yaml
project: my_feature_repo
registry: data/registry.db
provider: local
online_store:
type: milvus
host: "http://localhost"
port: 19530
username: "username"
password: "password"
embedding_dim: 128
index_type: "IVF_FLAT"
metric_type: "COSINE"
```
{% endcode %}

## Configuration options

| Option | Default | Description |
|:-------|:--------|:------------|
| `path` | `""` | Path to a Milvus Lite database file. Used when `provider: local` and `path` is set. |
| `host` | `http://localhost` | Milvus server host, including the scheme. |
| `port` | `19530` | Milvus server port. |
| `username` / `password` | `""` | Credentials, sent as the token `username:password`. |
| `embedding_dim` | `128` | Dimension of vector fields. |
| `index_type` | `FLAT` | Index type for vector fields with `vector_index=True`. |
| `metric_type` | `COSINE` | Default metric when a field does not set `vector_search_metric`. |
| `nlist` | `128` | `nlist` index parameter. |
| `vector_enabled` | `true` | Enables vector search. |
| `varchar_max_length` | `65535` | Default `max_length` of VARCHAR fields. Override per field with the `max_length` tag. |
| `enable_openai_compatible_store` | `false` | Store numeric features as native Milvus numeric types. |

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

## Collection loading

Feast creates collections together with their indexes, which makes Milvus load them straight away.
When Feast finds an existing collection it checks its load state and loads it only if needed.
Reads and searches never load collections, so a collection released outside Feast is only reloaded
the next time a Feast process first accesses it.

## Feature views without vectors

Milvus requires every collection to have a vector field. For feature views that have no vector
feature, Feast adds a 2-dimensional `_placeholder_vector` field with a FLAT index and fills it with zeros.
It is never returned or searched.

## Functionality Matrix

The set of functionality supported by online stores is described in detail [here](overview.md#functionality).
Expand Down
87 changes: 59 additions & 28 deletions sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
FieldSchema,
MilvusClient,
)
from pymilvus.client.types import LoadState

from feast import Entity
from feast.feature_view import FeatureView
Expand Down Expand Up @@ -114,6 +115,13 @@
DataType.BOOL,
}

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


def _milvus_escape_string(s: str) -> str:
"""Escape a string for safe use inside a Milvus single-quoted literal.
Expand Down Expand Up @@ -348,9 +356,9 @@ def _get_or_create_collection(
if not has_vector_field:
fields.append(
FieldSchema(
name="_placeholder_vector",
name=PLACEHOLDER_VECTOR_FIELD,
dtype=DataType.FLOAT_VECTOR,
dim=1,
dim=PLACEHOLDER_VECTOR_DIM,
)
)
schema = CollectionSchema(
Expand All @@ -360,22 +368,14 @@ def _get_or_create_collection(
collection_name=collection_name
)
if not collection_exists:
self.client.create_collection(
collection_name=collection_name,
dimension=config.online_store.embedding_dim,
schema=schema,
)
index_params = self.client.prepare_index_params()
indices_added = False
for vector_field in schema.fields:
if (
vector_field.dtype
in [
DataType.FLOAT_VECTOR,
DataType.BINARY_VECTOR,
]
and vector_field.name in vector_field_dict
):
if vector_field.dtype not in [
DataType.FLOAT_VECTOR,
DataType.BINARY_VECTOR,
]:
continue
if vector_field.name in vector_field_dict:
metric = vector_field_dict[
vector_field.name
].vector_search_metric
Expand All @@ -387,19 +387,44 @@ def _get_or_create_collection(
index_name=f"vector_index_{vector_field.name}",
params={"nlist": config.online_store.nlist},
)
indices_added = True
if indices_added:
self.client.create_index(
collection_name=collection_name,
index_params=index_params,
)
else:
# Vector fields that aren't searched (the placeholder,
# or arrays without vector_index) still need an index,
# otherwise Milvus servers refuse to load the collection.
index_params.add_index(
collection_name=collection_name,
field_name=vector_field.name,
metric_type="L2"
if vector_field.name == PLACEHOLDER_VECTOR_FIELD
else config.online_store.metric_type,
index_type="FLAT",
index_name=f"vector_index_{vector_field.name}",
)
# Every collection has at least one vector field, and every
# vector field is indexed, so passing the index params here
# makes Milvus create the indexes and load the collection.
self.client.create_collection(
collection_name=collection_name,
dimension=config.online_store.embedding_dim,
schema=schema,
index_params=index_params,
)
else:
self.client.load_collection(collection_name)
self._ensure_loaded(collection_name)
# Collections are only cached once loaded, so reads and searches
# don't need to load them again.
self._collections[collection_name] = self.client.describe_collection(
collection_name
)
return self._collections[collection_name]

def _ensure_loaded(self, collection_name: str) -> None:
"""Load an existing collection unless Milvus already has it loaded."""
assert self.client is not None, "Milvus client is not initialized"
load_state = self.client.get_load_state(collection_name).get("state")
if load_state != LoadState.Loaded:
self.client.load_collection(collection_name)

def online_write_batch(
self,
config: RepoConfig,
Expand All @@ -423,6 +448,15 @@ def online_write_batch(
collection_field_types = {
field["name"]: field["type"] for field in collection["fields"]
}
# Collections created by older Feast versions may use a 1-dim placeholder.
placeholder_dim = next(
(
int(field.get("params", {}).get("dim", PLACEHOLDER_VECTOR_DIM))
for field in collection["fields"]
if field["name"] == PLACEHOLDER_VECTOR_FIELD
),
PLACEHOLDER_VECTOR_DIM,
)
schema_internal_fields = {"event_ts", "created_ts"}
collection_has_native_numerics = any(
collection_field_types.get(name) in MILVUS_NATIVE_NUMERIC_TYPES
Expand Down Expand Up @@ -472,8 +506,8 @@ def online_write_batch(
for field in required_fields:
if field not in single_entity_record:
field_type = collection_field_types.get(field, DataType.VARCHAR)
if field == "_placeholder_vector":
single_entity_record[field] = [float("nan")]
if field == PLACEHOLDER_VECTOR_FIELD:
single_entity_record[field] = [0.0] * placeholder_dim
else:
single_entity_record[field] = _default_for_milvus_type(
field_type
Expand Down Expand Up @@ -533,7 +567,6 @@ def online_read(
+ ", ".join([f"'{e}'" for e in composite_entities])
+ "]"
)
self.client.load_collection(collection_name)
results = self.client.query(
collection_name=collection_name,
filter=query_filter_for_entities,
Expand Down Expand Up @@ -746,8 +779,6 @@ def retrieve_online_documents_v2(
ann_search_field = field["name"]
break

self.client.load_collection(collection_name)

if filters and filters_contain_numeric_comparison(filters):
collection_field_types = {
f["name"]: f["type"] for f in collection["fields"]
Expand Down
Loading