|
12 | 12 |
|
13 | 13 | from feast import Entity, FeatureView |
14 | 14 | from feast.field import Field |
| 15 | +from feast.filter_models import ComparisonFilter |
15 | 16 | from feast.infra.online_stores.milvus_online_store.milvus import ( |
16 | 17 | PLACEHOLDER_VECTOR_DIM, |
17 | 18 | PLACEHOLDER_VECTOR_FIELD, |
@@ -579,3 +580,169 @@ def test_collection_and_read_consistency_levels_independent() -> None: |
579 | 580 | def test_invalid_consistency_level_rejected(field: str) -> None: |
580 | 581 | with pytest.raises(ValidationError): |
581 | 582 | MilvusOnlineStoreConfig(**{field: "Immediate"}) |
| 583 | + |
| 584 | + |
| 585 | +def _catalog_feature_view(tags: Optional[Dict[str, str]] = None) -> FeatureView: |
| 586 | + return FeatureView( |
| 587 | + name="products", |
| 588 | + entities=[ |
| 589 | + Entity( |
| 590 | + name="product_id", join_keys=["product_id"], value_type=ValueType.INT64 |
| 591 | + ) |
| 592 | + ], |
| 593 | + ttl=timedelta(days=1), |
| 594 | + schema=[ |
| 595 | + Field(name="product_id", dtype=Int64), |
| 596 | + Field(name="brand_id", dtype=String), |
| 597 | + Field( |
| 598 | + name="embedding", |
| 599 | + dtype=Array(Float32), |
| 600 | + vector_index=True, |
| 601 | + vector_search_metric="COSINE", |
| 602 | + ), |
| 603 | + Field(name="title", dtype=String), |
| 604 | + Field(name="price", dtype=Float32), |
| 605 | + ], |
| 606 | + tags=tags or {}, |
| 607 | + ) |
| 608 | + |
| 609 | + |
| 610 | +def _product_key(product_id: int) -> EntityKeyProto: |
| 611 | + return EntityKeyProto( |
| 612 | + join_keys=["product_id"], entity_values=[ValueProto(int64_val=product_id)] |
| 613 | + ) |
| 614 | + |
| 615 | + |
| 616 | +def _write_products( |
| 617 | + store: MilvusOnlineStore, config: RepoConfig, fv: FeatureView |
| 618 | +) -> None: |
| 619 | + def embedding(x: float, y: float) -> ValueProto: |
| 620 | + value = ValueProto() |
| 621 | + value.float_list_val.val.extend([x, y]) |
| 622 | + return value |
| 623 | + |
| 624 | + now = datetime.now(timezone.utc) |
| 625 | + products = [ |
| 626 | + (1, "acme", (1.0, 0.0), "Acme kettle"), |
| 627 | + (2, "acme", (0.9, 0.1), "Acme toaster"), |
| 628 | + (3, "globex", (1.0, 0.0), "Globex kettle"), |
| 629 | + ] |
| 630 | + store.online_write_batch( |
| 631 | + config, |
| 632 | + fv, |
| 633 | + [ |
| 634 | + ( |
| 635 | + _product_key(product_id), |
| 636 | + { |
| 637 | + "brand_id": ValueProto(string_val=brand), |
| 638 | + "embedding": embedding(*vector), |
| 639 | + "title": ValueProto(string_val=title), |
| 640 | + }, |
| 641 | + now, |
| 642 | + now, |
| 643 | + ) |
| 644 | + for product_id, brand, vector, title in products |
| 645 | + ], |
| 646 | + progress=None, |
| 647 | + ) |
| 648 | + |
| 649 | + |
| 650 | +def _partition_key_field( |
| 651 | + store: MilvusOnlineStore, collection_name: str |
| 652 | +) -> Optional[str]: |
| 653 | + assert store.client is not None |
| 654 | + fields = store.client.describe_collection(collection_name)["fields"] |
| 655 | + return next((f["name"] for f in fields if f.get("is_partition_key")), None) |
| 656 | + |
| 657 | + |
| 658 | +@pytest.mark.parametrize( |
| 659 | + "online_store, tags", |
| 660 | + [ |
| 661 | + ({}, {"milvus.partition_key": "brand_id"}), |
| 662 | + ({"partition_key": "brand_id"}, {}), |
| 663 | + ], |
| 664 | + ids=["feature_view_tag", "store_config"], |
| 665 | +) |
| 666 | +def test_partition_key_filtering( |
| 667 | + tmp_path: Path, online_store: Dict[str, Any], tags: Dict[str, str] |
| 668 | +) -> None: |
| 669 | + config = _lite_config(tmp_path, **online_store) |
| 670 | + fv = _catalog_feature_view(tags) |
| 671 | + store = MilvusOnlineStore() |
| 672 | + store.update(config, [], [fv], [], [], partial=False) |
| 673 | + _write_products(store, config, fv) |
| 674 | + |
| 675 | + assert _partition_key_field(store, "test_milvus_products") == "brand_id" |
| 676 | + results = store.retrieve_online_documents_v2( |
| 677 | + config, |
| 678 | + fv, |
| 679 | + ["embedding", "brand_id", "title"], |
| 680 | + embedding=[1.0, 0.0], |
| 681 | + top_k=5, |
| 682 | + distance_metric="COSINE", |
| 683 | + filters=ComparisonFilter(type="eq", key="brand_id", value="acme"), |
| 684 | + ) |
| 685 | + |
| 686 | + titles = sorted(values["title"].string_val for _, _, values in results if values) |
| 687 | + assert titles == ["Acme kettle", "Acme toaster"] |
| 688 | + |
| 689 | + |
| 690 | +def test_store_partition_key_ignored_for_feature_views_without_the_field( |
| 691 | + tmp_path: Path, |
| 692 | +) -> None: |
| 693 | + config = _lite_config(tmp_path, partition_key="brand_id") |
| 694 | + fv = _scalar_feature_view() |
| 695 | + store = MilvusOnlineStore() |
| 696 | + store.update(config, [], [fv], [], [], partial=False) |
| 697 | + |
| 698 | + assert _partition_key_field(store, "test_milvus_driver_stats") is None |
| 699 | + |
| 700 | + |
| 701 | +def test_feature_view_tag_overrides_store_partition_key(tmp_path: Path) -> None: |
| 702 | + config = _lite_config(tmp_path, partition_key="brand_id") |
| 703 | + fv = _catalog_feature_view({"milvus.partition_key": "title"}) |
| 704 | + store = MilvusOnlineStore() |
| 705 | + store.update(config, [], [fv], [], [], partial=False) |
| 706 | + |
| 707 | + assert _partition_key_field(store, "test_milvus_products") == "title" |
| 708 | + |
| 709 | + |
| 710 | +@pytest.mark.parametrize( |
| 711 | + "online_store, partition_key, error", |
| 712 | + [ |
| 713 | + ({}, "missing_field", "is not a field"), |
| 714 | + # Native numeric storage makes price a FLOAT, which can't be a partition key. |
| 715 | + ({"enable_openai_compatible_store": True}, "price", "must be INT64 or VARCHAR"), |
| 716 | + ], |
| 717 | +) |
| 718 | +@patch(f"{MILVUS_MODULE}.MilvusClient") |
| 719 | +def test_invalid_partition_key( |
| 720 | + mock_client_cls: MagicMock, |
| 721 | + online_store: Dict[str, Any], |
| 722 | + partition_key: str, |
| 723 | + error: str, |
| 724 | +) -> None: |
| 725 | + mock_client = _mock_client(mock_client_cls, has_collection=False) |
| 726 | + fv = _catalog_feature_view({"milvus.partition_key": partition_key}) |
| 727 | + |
| 728 | + with pytest.raises(ValueError, match=error): |
| 729 | + MilvusOnlineStore()._get_or_create_collection(_mock_config(**online_store), fv) |
| 730 | + mock_client.create_collection.assert_not_called() |
| 731 | + |
| 732 | + |
| 733 | +@patch(f"{MILVUS_MODULE}.MilvusClient") |
| 734 | +def test_warns_when_existing_collection_lacks_partition_key( |
| 735 | + mock_client_cls: MagicMock, caplog: pytest.LogCaptureFixture |
| 736 | +) -> None: |
| 737 | + mock_client = _mock_client(mock_client_cls, has_collection=True) |
| 738 | + mock_client.get_load_state.return_value = {"state": LoadState.Loaded} |
| 739 | + mock_client.describe_collection.return_value = { |
| 740 | + "collection_name": "test_milvus_products", |
| 741 | + "fields": [{"name": "brand_id", "type": DataType.VARCHAR, "params": {}}], |
| 742 | + } |
| 743 | + fv = _catalog_feature_view({"milvus.partition_key": "brand_id"}) |
| 744 | + |
| 745 | + MilvusOnlineStore()._get_or_create_collection(_mock_config(), fv) |
| 746 | + |
| 747 | + assert "created without partition key 'brand_id'" in caplog.text |
| 748 | + mock_client.create_collection.assert_not_called() |
0 commit comments