diff --git a/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py b/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py index 29a6edf30ad..b78c78ee0eb 100644 --- a/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py +++ b/sdk/python/feast/infra/online_stores/qdrant_online_store/qdrant.py @@ -14,6 +14,7 @@ get_list_val_str, serialize_entity_key, ) +from feast.infra.online_stores.helpers import compute_versioned_name from feast.infra.online_stores.online_store import OnlineStore from feast.infra.online_stores.vector_store import VectorStoreConfig from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto @@ -27,6 +28,21 @@ SCROLL_SIZE = 1000 + +def _collection_name(config: RepoConfig, table: FeatureView) -> str: + """Qdrant collection backing a feature view, suffixed with ``_v{N}`` when + online feature view versioning is enabled. + + Deliberately built from ``compute_versioned_name`` rather than + ``compute_table_id``: Qdrant collections have always been named by the bare + feature view name, so adding the ``{project}_`` prefix that other stores use + would orphan every existing collection. + """ + return compute_versioned_name( + table, config.registry.enable_online_feature_view_versioning + ) + + DISTANCE_MAPPING = { "cosine": models.Distance.COSINE, "l2": models.Distance.EUCLID, @@ -139,7 +155,7 @@ def online_write_batch( ) self._get_client(config).upload_points( - collection_name=table.name, + collection_name=_collection_name(config, table), batch_size=config.online_store.write_batch_size, points=points, wait=True, @@ -172,7 +188,7 @@ def online_read( stop_scrolling = False while not stop_scrolling: records, next_offset = self._get_client(config).scroll( - collection_name=config.online_store.collection_name, + collection_name=_collection_name(config, table), limit=SCROLL_SIZE, offset=next_offset, with_payload=True, @@ -209,7 +225,7 @@ def create_collection(self, config: RepoConfig, table: FeatureView): client: QdrantClient = self._get_client(config) client.create_collection( - collection_name=table.name, + collection_name=_collection_name(config, table), vectors_config={ config.online_store.vector_name: models.VectorParams( size=vector_field_length, @@ -218,12 +234,12 @@ def create_collection(self, config: RepoConfig, table: FeatureView): }, ) client.create_payload_index( - collection_name=table.name, + collection_name=_collection_name(config, table), field_name="entity_key", field_schema=models.PayloadSchemaType.KEYWORD, ) client.create_payload_index( - collection_name=table.name, + collection_name=_collection_name(config, table), field_name="feature_name", field_schema=models.PayloadSchemaType.KEYWORD, ) @@ -238,7 +254,9 @@ def update( partial: bool, ): for table in tables_to_delete: - self._get_client(config).delete_collection(collection_name=table.name) + self._get_client(config).delete_collection( + collection_name=_collection_name(config, table) + ) for table in tables_to_keep: self.create_collection(config, table) @@ -251,7 +269,9 @@ def teardown( project = config.project try: for table in tables: - self._get_client(config).delete_collection(collection_name=table.name) + self._get_client(config).delete_collection( + collection_name=_collection_name(config, table) + ) except Exception as e: logging.exception(f"Error deleting collection in project {project}: {e}") raise @@ -288,7 +308,7 @@ def retrieve_online_documents( points = ( self._get_client(config) .query_points( - collection_name=table.name, + collection_name=_collection_name(config, table), query=embedding, limit=top_k, with_payload=True, diff --git a/sdk/python/tests/unit/infra/online_store/test_qdrant_online_store.py b/sdk/python/tests/unit/infra/online_store/test_qdrant_online_store.py new file mode 100644 index 00000000000..885ba5d6c58 --- /dev/null +++ b/sdk/python/tests/unit/infra/online_store/test_qdrant_online_store.py @@ -0,0 +1,137 @@ +from datetime import datetime, timedelta + +import pytest + +from feast import Entity, FeatureView, Field, FileSource, RepoConfig +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto +from feast.types import Int64, String +from feast.value_type import ValueType + +qdrant = pytest.importorskip("qdrant_client") + +from feast.infra.online_stores.qdrant_online_store.qdrant import ( # noqa: E402 + QdrantOnlineStore, + QdrantOnlineStoreConfig, + _collection_name, +) + + +def _config(versioning: bool) -> RepoConfig: + return RepoConfig( + project="test_project", + online_store=QdrantOnlineStoreConfig( + type="qdrant", + location=":memory:", + vector_enabled=True, + similarity="cosine", + ), + registry={ + "path": "dummy_registry", + "enable_online_feature_view_versioning": versioning, + }, + entity_key_serialization_version=3, + ) + + +def _feature_view(version: int = 0) -> FeatureView: + fv = FeatureView( + name="driver_stats", + entities=[Entity(name="driver_id", value_type=ValueType.INT64)], + ttl=timedelta(days=1), + schema=[ + Field(name="driver_id", dtype=Int64), + Field(name="trips", dtype=String), + ], + source=FileSource(path="driver.parquet", timestamp_field="event_timestamp"), + ) + if version: + fv.projection.version_tag = version + return fv + + +def _write(store, config, fv, value): + entity_key = EntityKeyProto( + join_keys=["driver_id"], entity_values=[ValueProto(int64_val=1)] + ) + store.online_write_batch( + config, + fv, + [ + ( + entity_key, + {"trips": ValueProto(string_val=value)}, + datetime(2024, 1, 1), + None, + ) + ], + None, + ) + + +class TestQdrantCollectionNaming: + def test_unversioned_name_is_unchanged(self): + """Existing deployments must keep their current collection names.""" + assert _collection_name(_config(False), _feature_view()) == "driver_stats" + + def test_versioning_disabled_ignores_the_version_tag(self): + assert ( + _collection_name(_config(False), _feature_view(version=2)) == "driver_stats" + ) + + def test_versioned_name_gets_a_suffix(self): + assert ( + _collection_name(_config(True), _feature_view(version=2)) + == "driver_stats_v2" + ) + + +class TestQdrantVersionedCollections: + """End-to-end against an in-process Qdrant.""" + + def test_versions_write_to_separate_collections(self): + config = _config(True) + store = QdrantOnlineStore() + v1, v2 = _feature_view(version=1), _feature_view(version=2) + + store.update(config, [], [v1, v2], [], [], partial=False) + names = { + c.name for c in store._get_client(config).get_collections().collections + } + assert {"driver_stats_v1", "driver_stats_v2"} <= names + + _write(store, config, v1, "from_v1") + _write(store, config, v2, "from_v2") + + client = store._get_client(config) + assert client.count("driver_stats_v1").count == 1 + assert client.count("driver_stats_v2").count == 1 + + payload = client.scroll("driver_stats_v1", limit=10, with_payload=True)[0] + assert [p.payload["feature_name"] for p in payload] == ["trips"] + + def test_teardown_removes_only_the_targeted_version(self): + config = _config(True) + store = QdrantOnlineStore() + v1, v2 = _feature_view(version=1), _feature_view(version=2) + store.update(config, [], [v1, v2], [], [], partial=False) + + store.teardown(config, [v1], []) + + names = { + c.name for c in store._get_client(config).get_collections().collections + } + assert "driver_stats_v1" not in names + assert "driver_stats_v2" in names + + def test_unversioned_store_still_round_trips(self): + """The default path must be untouched by the versioning change.""" + config = _config(False) + store = QdrantOnlineStore() + fv = _feature_view() + + store.update(config, [], [fv], [], [], partial=False) + _write(store, config, fv, "value") + + client = store._get_client(config) + assert client.count("driver_stats").count == 1