Visitar URL original
feat: Support Milvus partition keys in the Milvus online store by simonhearne · Pull Request #6940 · feast-dev/feast · GitHub
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
46 changes: 46 additions & 0 deletions docs/reference/online-stores/milvus.md
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ online_store:
| `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`. |
| `consistency_level` | unset | `Strong`, `Bounded`, `Session` or `Eventually`. Sent with every read and search. When unset, Milvus uses the collection's level. |
| `collection_consistency_level` | unset | `Strong`, `Bounded`, `Session` or `Eventually`. Set when Feast creates a collection. When unset, Milvus uses its default (`Bounded`). |
| `partition_key` | unset | Field to use as the Milvus partition key in feature views that contain it. See [Partition key](#partition-key). |
| `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. |
Expand Down Expand Up @@ -132,6 +133,51 @@ online_store:
Index parameters only apply when Feast creates a collection. To change them for an existing
collection, run `feast teardown` and `feast apply`, then materialize again.

## Partition key

A [partition key](https://milvus.io/docs/use-partition-key.md) makes Milvus group rows by the key's
value, so searches filtered on it only scan the matching partitions. This suits multi-tenant data,
such as a catalogue shared by many brands.

Set the partition key per feature view with the `milvus.partition_key` tag:

```python
products = FeatureView(
name="products",
entities=[product],
schema=[
Field(name="product_id", dtype=Int64),
Field(name="brand_id", dtype=String),
Field(name="embedding", dtype=Array(Float32), vector_index=True),
Field(name="title", dtype=String),
],
source=products_source,
tags={"milvus.partition_key": "brand_id"},
)
```

or for every feature view that has the field, with `partition_key: brand_id` in the online store
config. The tag takes precedence. Then filter on the key when retrieving:

```python
store.retrieve_online_documents_v2(
features=["products:embedding", "products:title"],
query=query_embedding,
top_k=10,
filters=ComparisonFilter(type="eq", key="brand_id", value="acme"),
)
```

The partition key field must be stored as `VARCHAR` or `INT64`. String fields are always `VARCHAR`,
and integer fields are `VARCHAR` unless native numeric types are enabled, in which case `Int64` works
and `Int32` does not.

{% hint style="warning" %}
The partition key only applies when Feast creates a collection. Existing collections are not
changed; Feast logs a warning for them. To add a partition key to an existing feature view, run
`feast teardown` and `feast apply`, then materialize again.
{% endhint %}

## Consistency level

By default Milvus uses `Bounded` consistency, so a read issued straight after materialization may
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,9 @@ class MilvusOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig):
collection_consistency_level: Optional[
Literal["Strong", "Bounded", "Session", "Eventually"]
] = None
# Field to use as the Milvus partition key in feature views that contain it.
# A feature view's "milvus.partition_key" tag takes precedence.
partition_key: Optional[StrictStr] = None
username: Optional[StrictStr] = ""
password: Optional[StrictStr] = ""
enable_openai_compatible_store: Optional[bool] = False
Expand Down Expand Up @@ -386,6 +389,9 @@ def _get_or_create_collection(
dim=PLACEHOLDER_VECTOR_DIM,
)
)
partition_key = _partition_key_name(config.online_store, table)
if partition_key:
_mark_partition_key(fields, partition_key, table.name)
schema = CollectionSchema(
fields=fields, description="Feast feature view data"
)
Expand Down Expand Up @@ -440,6 +446,10 @@ def _get_or_create_collection(
self._collections[collection_name] = self.client.describe_collection(
collection_name
)
if collection_exists and partition_key:
_warn_if_partition_key_missing(
self._collections[collection_name], partition_key
)
return self._collections[collection_name]

def _ensure_loaded(self, collection_name: str) -> None:
Expand Down Expand Up @@ -1068,6 +1078,63 @@ def _search_params(online_config: MilvusOnlineStoreConfig) -> Dict[str, Any]:
return {"nprobe": 10}


PARTITION_KEY_TAG = "milvus.partition_key"
PARTITION_KEY_TYPES = {DataType.INT64, DataType.VARCHAR}


def _partition_key_name(
online_config: MilvusOnlineStoreConfig, table: FeatureView
) -> Optional[str]:
"""Return the partition key field for a feature view, if one is configured.

The feature view's ``milvus.partition_key`` tag takes precedence over the
store-level ``partition_key``, which only applies to feature views that
contain that field.
"""
tagged = table.tags.get(PARTITION_KEY_TAG)
if tagged:
return tagged
configured = online_config.partition_key
if configured and any(field.name == configured for field in table.schema):
return configured
return None


def _mark_partition_key(
fields: List[FieldSchema], partition_key: str, table_name: str
) -> None:
field = next((f for f in fields if f.name == partition_key), None)
if field is None:
raise ValueError(
f"Partition key '{partition_key}' is not a field of feature view "
f"'{table_name}'."
)
if field.dtype not in PARTITION_KEY_TYPES:
raise ValueError(
f"Partition key '{partition_key}' of feature view '{table_name}' is "
f"stored as {field.dtype.name}, but Milvus partition keys must be "
"INT64 or VARCHAR."
)
field.is_partition_key = True


def _warn_if_partition_key_missing(
collection: Dict[str, Any], partition_key: str
) -> None:
has_key = any(
field["name"] == partition_key and field.get("is_partition_key")
for field in collection["fields"]
)
if not has_key:
logger.warning(
"Collection '%s' was created without partition key '%s'. The partition "
"key only applies when a collection is created: run `feast teardown` "
"and `feast apply`, then materialize again to use it.",
collection["collection_name"],
partition_key,
)


def _consistency_kwargs(online_config: MilvusOnlineStoreConfig) -> Dict[str, Any]:
"""Read and search kwargs; only pass a level when configured."""
if online_config.consistency_level:
Expand Down
47 changes: 47 additions & 0 deletions sdk/python/tests/integration/online_store/test_milvus_remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@

from feast import Entity, FeatureView
from feast.field import Field
from feast.filter_models import ComparisonFilter
from feast.infra.online_stores.milvus_online_store.milvus import MilvusOnlineStore
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
Expand Down Expand Up @@ -360,3 +361,49 @@ def test_consistency_level(
)
rows = _read(store, config, fv, [1], ["city"])
assert rows[0] is not None and rows[0]["city"].string_val == "Oslo"


def test_partition_key_filtering(
tmp_path: Path, project: str, store: MilvusOnlineStore
) -> None:
config = _repo_config(tmp_path, project, consistency_level="Strong")
fv = FeatureView(
name="products",
entities=[
Entity(
name="driver_id", join_keys=["driver_id"], value_type=ValueType.INT64
)
],
ttl=timedelta(days=1),
schema=[
Field(name="driver_id", dtype=Int64),
Field(name="brand_id", dtype=String),
Field(
name="embedding",
dtype=Array(Float32),
vector_index=True,
vector_search_metric="COSINE",
),
Field(name="city", dtype=String),
],
tags={"milvus.partition_key": "brand_id"},
)
store.update(config, [], [fv], [], [], partial=False)
rows = _vector_rows()
rows[1]["brand_id"] = ValueProto(string_val="acme")
rows[2]["brand_id"] = ValueProto(string_val="globex")
_write_rows(store, config, fv, rows)

assert store.client is not None
fields = store.client.describe_collection(f"{project}_{fv.name}")["fields"]
assert [f["name"] for f in fields if f.get("is_partition_key")] == ["brand_id"]

hits = _search(
store,
config,
fv,
[1.0, 0.0],
top_k=5,
filters=ComparisonFilter(type="eq", key="brand_id", value="globex"),
)
assert [hit["city"].string_val for hit in hits] == ["Rome"]
167 changes: 167 additions & 0 deletions sdk/python/tests/unit/infra/online_store/test_milvus_online_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@

from feast import Entity, FeatureView
from feast.field import Field
from feast.filter_models import ComparisonFilter
from feast.infra.online_stores.milvus_online_store.milvus import (
PLACEHOLDER_VECTOR_DIM,
PLACEHOLDER_VECTOR_FIELD,
Expand Down Expand Up @@ -579,3 +580,169 @@ def test_collection_and_read_consistency_levels_independent() -> None:
def test_invalid_consistency_level_rejected(field: str) -> None:
with pytest.raises(ValidationError):
MilvusOnlineStoreConfig(**{field: "Immediate"})


def _catalog_feature_view(tags: Optional[Dict[str, str]] = None) -> FeatureView:
return FeatureView(
name="products",
entities=[
Entity(
name="product_id", join_keys=["product_id"], value_type=ValueType.INT64
)
],
ttl=timedelta(days=1),
schema=[
Field(name="product_id", dtype=Int64),
Field(name="brand_id", dtype=String),
Field(
name="embedding",
dtype=Array(Float32),
vector_index=True,
vector_search_metric="COSINE",
),
Field(name="title", dtype=String),
Field(name="price", dtype=Float32),
],
tags=tags or {},
)


def _product_key(product_id: int) -> EntityKeyProto:
return EntityKeyProto(
join_keys=["product_id"], entity_values=[ValueProto(int64_val=product_id)]
)


def _write_products(
store: MilvusOnlineStore, config: RepoConfig, fv: FeatureView
) -> None:
def embedding(x: float, y: float) -> ValueProto:
value = ValueProto()
value.float_list_val.val.extend([x, y])
return value

now = datetime.now(timezone.utc)
products = [
(1, "acme", (1.0, 0.0), "Acme kettle"),
(2, "acme", (0.9, 0.1), "Acme toaster"),
(3, "globex", (1.0, 0.0), "Globex kettle"),
]
store.online_write_batch(
config,
fv,
[
(
_product_key(product_id),
{
"brand_id": ValueProto(string_val=brand),
"embedding": embedding(*vector),
"title": ValueProto(string_val=title),
},
now,
now,
)
for product_id, brand, vector, title in products
],
progress=None,
)


def _partition_key_field(
store: MilvusOnlineStore, collection_name: str
) -> Optional[str]:
assert store.client is not None
fields = store.client.describe_collection(collection_name)["fields"]
return next((f["name"] for f in fields if f.get("is_partition_key")), None)


@pytest.mark.parametrize(
"online_store, tags",
[
({}, {"milvus.partition_key": "brand_id"}),
({"partition_key": "brand_id"}, {}),
],
ids=["feature_view_tag", "store_config"],
)
def test_partition_key_filtering(
tmp_path: Path, online_store: Dict[str, Any], tags: Dict[str, str]
) -> None:
config = _lite_config(tmp_path, **online_store)
fv = _catalog_feature_view(tags)
store = MilvusOnlineStore()
store.update(config, [], [fv], [], [], partial=False)
_write_products(store, config, fv)

assert _partition_key_field(store, "test_milvus_products") == "brand_id"
results = store.retrieve_online_documents_v2(
config,
fv,
["embedding", "brand_id", "title"],
embedding=[1.0, 0.0],
top_k=5,
distance_metric="COSINE",
filters=ComparisonFilter(type="eq", key="brand_id", value="acme"),
)

titles = sorted(values["title"].string_val for _, _, values in results if values)
assert titles == ["Acme kettle", "Acme toaster"]


def test_store_partition_key_ignored_for_feature_views_without_the_field(
tmp_path: Path,
) -> None:
config = _lite_config(tmp_path, partition_key="brand_id")
fv = _scalar_feature_view()
store = MilvusOnlineStore()
store.update(config, [], [fv], [], [], partial=False)

assert _partition_key_field(store, "test_milvus_driver_stats") is None


def test_feature_view_tag_overrides_store_partition_key(tmp_path: Path) -> None:
config = _lite_config(tmp_path, partition_key="brand_id")
fv = _catalog_feature_view({"milvus.partition_key": "title"})
store = MilvusOnlineStore()
store.update(config, [], [fv], [], [], partial=False)

assert _partition_key_field(store, "test_milvus_products") == "title"


@pytest.mark.parametrize(
"online_store, partition_key, error",
[
({}, "missing_field", "is not a field"),
# Native numeric storage makes price a FLOAT, which can't be a partition key.
({"enable_openai_compatible_store": True}, "price", "must be INT64 or VARCHAR"),
],
)
@patch(f"{MILVUS_MODULE}.MilvusClient")
def test_invalid_partition_key(
mock_client_cls: MagicMock,
online_store: Dict[str, Any],
partition_key: str,
error: str,
) -> None:
mock_client = _mock_client(mock_client_cls, has_collection=False)
fv = _catalog_feature_view({"milvus.partition_key": partition_key})

with pytest.raises(ValueError, match=error):
MilvusOnlineStore()._get_or_create_collection(_mock_config(**online_store), fv)
mock_client.create_collection.assert_not_called()


@patch(f"{MILVUS_MODULE}.MilvusClient")
def test_warns_when_existing_collection_lacks_partition_key(
mock_client_cls: MagicMock, caplog: pytest.LogCaptureFixture
) -> None:
mock_client = _mock_client(mock_client_cls, has_collection=True)
mock_client.get_load_state.return_value = {"state": LoadState.Loaded}
mock_client.describe_collection.return_value = {
"collection_name": "test_milvus_products",
"fields": [{"name": "brand_id", "type": DataType.VARCHAR, "params": {}}],
}
fv = _catalog_feature_view({"milvus.partition_key": "brand_id"})

MilvusOnlineStore()._get_or_create_collection(_mock_config(), fv)

assert "created without partition key 'brand_id'" in caplog.text
mock_client.create_collection.assert_not_called()
Loading