@@ -1714,3 +1714,181 @@ def test_milvus_keyword_search() -> None:
17141714 assert len (result_hybrid ["content" ]) > 0
17151715 assert any ("Feast" in content for content in result_hybrid ["content" ])
17161716 assert len (result_hybrid ["vector" ]) > 0
1717+
1718+
1719+ def test_milvus_update_preserves_collection_cache () -> None :
1720+ """
1721+ Regression test: update() used to overwrite self._collections with the
1722+ describe_collection() dict of the last processed table, replacing the
1723+ dict-of-dicts cache with a single flat dict. After the fix, each call
1724+ to _get_or_create_collection() updates the keyed entry in-place and the
1725+ cache remains a proper mapping from collection name to collection info.
1726+ """
1727+ from datetime import timedelta
1728+
1729+ from feast import Entity , FeatureView , Field , FileSource
1730+ from feast .types import Array , Float32 , Int64 , String
1731+
1732+ runner = CliRunner ()
1733+ with runner .local_repo (
1734+ example_repo_py = get_example_repo ("example_rag_feature_repo.py" ),
1735+ offline_store = "file" ,
1736+ online_store = "milvus" ,
1737+ apply = False ,
1738+ teardown = False ,
1739+ ) as store :
1740+ source = FileSource (
1741+ path = "data/dummy.parquet" ,
1742+ timestamp_field = "event_timestamp" ,
1743+ created_timestamp_column = "created_timestamp" ,
1744+ )
1745+ entity_a = Entity (name = "id_a" , join_keys = ["id_a" ], value_type = ValueType .INT64 )
1746+ entity_b = Entity (name = "id_b" , join_keys = ["id_b" ], value_type = ValueType .INT64 )
1747+
1748+ fv_a = FeatureView (
1749+ name = "fv_a" ,
1750+ entities = [entity_a ],
1751+ schema = [
1752+ Field (name = "id_a" , dtype = Int64 ),
1753+ Field (
1754+ name = "vec_a" ,
1755+ dtype = Array (Float32 ),
1756+ vector_index = True ,
1757+ vector_search_metric = "COSINE" ,
1758+ ),
1759+ Field (name = "text_a" , dtype = String ),
1760+ ],
1761+ source = source ,
1762+ ttl = timedelta (hours = 1 ),
1763+ )
1764+ fv_b = FeatureView (
1765+ name = "fv_b" ,
1766+ entities = [entity_b ],
1767+ schema = [
1768+ Field (name = "id_b" , dtype = Int64 ),
1769+ Field (
1770+ name = "vec_b" ,
1771+ dtype = Array (Float32 ),
1772+ vector_index = True ,
1773+ vector_search_metric = "COSINE" ,
1774+ ),
1775+ Field (name = "text_b" , dtype = String ),
1776+ ],
1777+ source = source ,
1778+ ttl = timedelta (hours = 1 ),
1779+ )
1780+
1781+ store .apply ([source , entity_a , entity_b , fv_a , fv_b ])
1782+
1783+ online_store = store ._provider ._online_store
1784+ # After applying two feature views, the cache must be a proper dict
1785+ # mapping collection names to collection-info dicts, not a flat dict.
1786+ assert isinstance (online_store ._collections , dict ), (
1787+ "_collections should be a dict"
1788+ )
1789+ collection_name_a = f"{ store .config .project }
1790+ collection_name_b = f"{ store .config .project }
1791+ assert collection_name_a in online_store ._collections , (
1792+ f"Cache missing entry for { collection_name_a }
1793+ )
1794+ assert collection_name_b in online_store ._collections , (
1795+ f"Cache missing entry for { collection_name_b }
1796+ "update() likely overwrote _collections with a single collection dict"
1797+ )
1798+ # Each cached value must be a collection-info dict (has a 'fields' key),
1799+ # not itself keyed by collection name.
1800+ for name in [collection_name_a , collection_name_b ]:
1801+ assert "fields" in online_store ._collections [name ], (
1802+ f"Cache entry for { name }
1803+ )
1804+
1805+
1806+ def test_milvus_plan_returns_empty_list () -> None :
1807+ """
1808+ Regression test: plan() used to raise NotImplementedError, causing
1809+ `feast plan` to crash for any project using the Milvus online store.
1810+ It should return [] matching the OnlineStore base class default.
1811+ """
1812+ from feast .infra .online_stores .milvus_online_store .milvus import MilvusOnlineStore
1813+
1814+ store = MilvusOnlineStore ()
1815+ result = store .plan (config = None , desired_registry_proto = None ) # type: ignore[arg-type]
1816+ assert result == [], f"plan() should return [] but returned { result !r}
1817+
1818+
1819+ def test_milvus_retrieve_online_documents_v2_missing_entity_key () -> None :
1820+ """
1821+ Regression test: retrieve_online_documents_v2() passed the raw
1822+ hit.get("entity", {}).get(composite_key_name, None) directly to
1823+ bytes.fromhex(), raising TypeError when the key was absent.
1824+ After the fix, a missing composite key produces a None entity_key_proto
1825+ instead of crashing.
1826+ """
1827+ from datetime import timedelta
1828+ from unittest .mock import patch
1829+
1830+ from feast import Entity , FeatureView , Field , FileSource
1831+ from feast .types import Array , Float32 , Int64 , String
1832+
1833+ runner = CliRunner ()
1834+ with runner .local_repo (
1835+ example_repo_py = get_example_repo ("example_rag_feature_repo.py" ),
1836+ offline_store = "file" ,
1837+ online_store = "milvus" ,
1838+ apply = False ,
1839+ teardown = False ,
1840+ ) as store :
1841+ source = FileSource (
1842+ path = "data/dummy.parquet" ,
1843+ timestamp_field = "event_timestamp" ,
1844+ created_timestamp_column = "created_timestamp" ,
1845+ )
1846+ entity = Entity (name = "doc_id" , join_keys = ["doc_id" ], value_type = ValueType .INT64 )
1847+ fv = FeatureView (
1848+ name = "docs" ,
1849+ entities = [entity ],
1850+ schema = [
1851+ Field (name = "doc_id" , dtype = Int64 ),
1852+ Field (
1853+ name = "vec" ,
1854+ dtype = Array (Float32 ),
1855+ vector_index = True ,
1856+ vector_search_metric = "COSINE" ,
1857+ ),
1858+ Field (name = "text" , dtype = String ),
1859+ ],
1860+ source = source ,
1861+ ttl = timedelta (hours = 1 ),
1862+ )
1863+ store .apply ([source , entity , fv ])
1864+
1865+ online_store = store ._provider ._online_store
1866+ fv_obj = store .get_feature_view ("docs" )
1867+ # Simulate a search hit that is missing the composite primary key.
1868+ fake_hit = {
1869+ "entity" : {
1870+ "event_ts" : int (_utc_now ().timestamp () * 1e6 ),
1871+ "created_ts" : int (_utc_now ().timestamp () * 1e6 ),
1872+ "text" : "hello" ,
1873+ },
1874+ "distance" : 0.9 ,
1875+ }
1876+
1877+ mock_results = [[fake_hit ]]
1878+ with patch .object (online_store .client , "search" , return_value = mock_results ):
1879+ with patch .object (
1880+ online_store .client , "load_collection" , return_value = None
1881+ ):
1882+ # Before the fix this raised TypeError: fromhex argument must be str, not None
1883+ result = online_store .retrieve_online_documents_v2 (
1884+ config = store .config ,
1885+ table = fv_obj ,
1886+ requested_features = ["text" ],
1887+ embedding = [0.1 ] * 10 ,
1888+ top_k = 1 ,
1889+ )
1890+ assert len (result ) == 1
1891+ _ts , entity_key_proto , _features = result [0 ]
1892+ assert entity_key_proto is None , (
1893+ "entity_key_proto should be None when the composite key is absent from the hit"
1894+ )
0 commit comments