diff --git a/‎docs/reference/online-stores/milvus.md‎
Lines changed: 46 additions & 0 deletions b/‎docs/reference/online-stores/milvus.md‎
Lines changed: 46 additions & 0 deletions
diff --git a/‎sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py‎
Lines changed: 67 additions & 0 deletions b/‎sdk/python/feast/infra/online_stores/milvus_online_store/milvus.py‎
Lines changed: 67 additions & 0 deletions
diff --git a/‎sdk/python/tests/integration/online_store/test_milvus_remote.py‎
Lines changed: 47 additions & 0 deletions b/‎sdk/python/tests/integration/online_store/test_milvus_remote.py‎
Lines changed: 47 additions & 0 deletions
diff --git a/‎sdk/python/tests/unit/infra/online_store/test_milvus_online_store.py‎
Lines changed: 167 additions & 0 deletions b/‎sdk/python/tests/unit/infra/online_store/test_milvus_online_store.py‎
Lines changed: 167 additions & 0 deletions
@@ -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. |
@@ -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
 
@@ -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
@@ -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"
             )
@@ -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:
@@ -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:
 
@@ -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
@@ -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"]
@@ -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,
@@ -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()