Skip to content
Open
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
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
)
Expand All @@ -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)

Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand Down
137 changes: 137 additions & 0 deletions sdk/python/tests/unit/infra/online_store/test_qdrant_online_store.py
Original file line number Diff line number Diff line change
@@ -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