|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
| 3 | +import time |
3 | 4 | from datetime import datetime |
4 | 5 | from logging import getLogger |
5 | | -from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple |
| 6 | +from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union |
6 | 7 |
|
7 | 8 | try: |
8 | 9 | from pymongo import AsyncMongoClient, MongoClient, UpdateOne |
9 | 10 | from pymongo.asynchronous.collection import AsyncCollection |
10 | 11 | from pymongo.collection import Collection |
11 | 12 | from pymongo.driver_info import DriverInfo |
| 13 | + from pymongo.operations import SearchIndexModel |
12 | 14 | except ImportError as e: |
13 | 15 | from feast.errors import FeastExtrasDependencyImportError |
14 | 16 |
|
15 | 17 | raise FeastExtrasDependencyImportError("mongodb", str(e)) |
16 | 18 |
|
17 | 19 | import feast.version |
| 20 | +from feast.batch_feature_view import BatchFeatureView |
18 | 21 | from feast.entity import Entity |
19 | 22 | from feast.feature_view import FeatureView |
20 | | -from feast.infra.key_encoding_utils import serialize_entity_key |
| 23 | +from feast.infra.key_encoding_utils import deserialize_entity_key, serialize_entity_key |
21 | 24 | from feast.infra.online_stores.online_store import OnlineStore |
| 25 | +from feast.infra.online_stores.vector_store import VectorStoreConfig |
22 | 26 | from feast.infra.supported_async_methods import SupportedAsyncMethods |
23 | 27 | from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto |
24 | 28 | from feast.protos.feast.types.Value_pb2 import Value as ValueProto |
25 | 29 | from feast.repo_config import FeastConfigBaseModel, RepoConfig |
| 30 | +from feast.stream_feature_view import StreamFeatureView |
26 | 31 | from feast.type_map import ( |
27 | 32 | feast_value_type_to_python_type, |
28 | 33 | python_values_to_proto_values, |
|
33 | 38 | DRIVER_METADATA = DriverInfo(name="Feast", version=feast.version.get_version()) |
34 | 39 |
|
35 | 40 |
|
36 | | -class MongoDBOnlineStoreConfig(FeastConfigBaseModel): |
| 41 | +class MongoDBOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig): |
37 | 42 | """MongoDB configuration. |
38 | 43 |
|
39 | 44 | For a description of kwargs that may be passed to MongoClient, |
@@ -48,6 +53,11 @@ class MongoDBOnlineStoreConfig(FeastConfigBaseModel): |
48 | 53 | ) |
49 | 54 | collection_suffix: str = "latest" |
50 | 55 | client_kwargs: Dict[str, Any] = {} |
| 56 | + # vector_enabled and similarity are inherited from VectorStoreConfig |
| 57 | + vector_index_wait_timeout: int = 60 |
| 58 | + """Seconds to wait for a newly created Atlas Search index to become READY.""" |
| 59 | + vector_index_wait_poll_interval: float = 1.0 |
| 60 | + """Seconds between polls when waiting for an Atlas Search index to become READY.""" |
51 | 61 |
|
52 | 62 |
|
53 | 63 | class MongoDBOnlineStore(OnlineStore): |
@@ -213,11 +223,135 @@ def online_read( |
213 | 223 |
|
214 | 224 | return self._convert_raw_docs_to_proto(ids, docs, table) |
215 | 225 |
|
| 226 | + def retrieve_online_documents_v2( |
| 227 | + self, |
| 228 | + config: RepoConfig, |
| 229 | + table: FeatureView, |
| 230 | + requested_features: List[str], |
| 231 | + embedding: Optional[List[float]], |
| 232 | + top_k: int, |
| 233 | + distance_metric: Optional[str] = None, |
| 234 | + query_string: Optional[str] = None, |
| 235 | + include_feature_view_version_metadata: bool = False, |
| 236 | + ) -> List[ |
| 237 | + Tuple[ |
| 238 | + Optional[datetime], |
| 239 | + Optional[EntityKeyProto], |
| 240 | + Optional[Dict[str, ValueProto]], |
| 241 | + ] |
| 242 | + ]: |
| 243 | + """Retrieve documents via MongoDB Atlas Vector Search ($vectorSearch). |
| 244 | +
|
| 245 | + Uses the ``$vectorSearch`` aggregation stage to find the *top_k* |
| 246 | + documents closest to the provided *embedding* vector. The method |
| 247 | + expects that an Atlas vector search index has already been created for |
| 248 | + the relevant field (see ``update()``). |
| 249 | +
|
| 250 | + Returns a list of 3-tuples ``(event_timestamp, entity_key_proto, |
| 251 | + feature_dict)`` where *feature_dict* includes the requested feature |
| 252 | + values plus a synthetic ``distance`` key with the vector search score. |
| 253 | + """ |
| 254 | + if not config.online_store.vector_enabled: |
| 255 | + raise ValueError( |
| 256 | + "Vector search is not enabled in the online store config. " |
| 257 | + "Set vector_enabled=True in MongoDBOnlineStoreConfig." |
| 258 | + ) |
| 259 | + if embedding is None: |
| 260 | + raise ValueError( |
| 261 | + "An embedding vector must be provided for MongoDB vector search." |
| 262 | + ) |
| 263 | + |
| 264 | + clxn = self._get_collection(config) |
| 265 | + |
| 266 | + # Identify the vector field on this feature view |
| 267 | + vector_fields = [f for f in table.features if f.vector_index] |
| 268 | + if not vector_fields: |
| 269 | + raise ValueError( |
| 270 | + f"Feature view '{table.name}' has no fields with vector_index=True." |
| 271 | + ) |
| 272 | + vector_field = vector_fields[0] |
| 273 | + path = f"features.{table.name}.{vector_field.name}" |
| 274 | + idx_name = self._vector_search_index_name(table.name, vector_field.name) |
| 275 | + |
| 276 | + # BSON cannot encode numpy float types — ensure native Python floats. |
| 277 | + query_vector = [float(v) for v in embedding] |
| 278 | + |
| 279 | + num_candidates = max(top_k * 10, 100) |
| 280 | + pipeline: List[Dict[str, Any]] = [ |
| 281 | + { |
| 282 | + "$vectorSearch": { |
| 283 | + "index": idx_name, |
| 284 | + "path": path, |
| 285 | + "queryVector": query_vector, |
| 286 | + "numCandidates": num_candidates, |
| 287 | + "limit": top_k, |
| 288 | + } |
| 289 | + }, |
| 290 | + { |
| 291 | + "$addFields": { |
| 292 | + "score": {"$meta": "vectorSearchScore"}, |
| 293 | + } |
| 294 | + }, |
| 295 | + ] |
| 296 | + |
| 297 | + results: List[ |
| 298 | + Tuple[ |
| 299 | + Optional[datetime], |
| 300 | + Optional[EntityKeyProto], |
| 301 | + Optional[Dict[str, ValueProto]], |
| 302 | + ] |
| 303 | + ] = [] |
| 304 | + |
| 305 | + for doc in clxn.aggregate(pipeline): |
| 306 | + # Deserialize entity key |
| 307 | + entity_key_bin = doc.get("_id") |
| 308 | + entity_key_proto = ( |
| 309 | + deserialize_entity_key( |
| 310 | + entity_key_bin, |
| 311 | + entity_key_serialization_version=config.entity_key_serialization_version, |
| 312 | + ) |
| 313 | + if entity_key_bin |
| 314 | + else None |
| 315 | + ) |
| 316 | + |
| 317 | + # Event timestamp |
| 318 | + event_ts = doc.get("event_timestamps", {}).get(table.name) |
| 319 | + |
| 320 | + # Build feature dict from raw doc values |
| 321 | + fv_features = doc.get("features", {}).get(table.name, {}) |
| 322 | + |
| 323 | + # Convert raw values → ValueProto for each requested feature |
| 324 | + feature_dict: Dict[str, ValueProto] = {} |
| 325 | + feature_type_map = {f.name: f.dtype.to_value_type() for f in table.features} |
| 326 | + |
| 327 | + for feat_name in requested_features: |
| 328 | + raw_val = fv_features.get(feat_name) |
| 329 | + if raw_val is not None: |
| 330 | + vtype = feature_type_map.get(feat_name) |
| 331 | + if vtype is not None: |
| 332 | + protos = python_values_to_proto_values( |
| 333 | + [raw_val], feature_type=vtype |
| 334 | + ) |
| 335 | + feature_dict[feat_name] = protos[0] |
| 336 | + else: |
| 337 | + # Fall back: try to store as-is |
| 338 | + feature_dict[feat_name] = _python_value_to_proto(raw_val) |
| 339 | + |
| 340 | + # Add distance (vector search score) |
| 341 | + score = doc.get("score", 0.0) |
| 342 | + feature_dict["distance"] = ValueProto(float_val=float(score)) |
| 343 | + |
| 344 | + results.append((event_ts, entity_key_proto, feature_dict)) |
| 345 | + |
| 346 | + return results |
| 347 | + |
216 | 348 | def update( |
217 | 349 | self, |
218 | 350 | config: RepoConfig, |
219 | 351 | tables_to_delete: Sequence[FeatureView], |
220 | | - tables_to_keep: Sequence[FeatureView], |
| 352 | + tables_to_keep: Sequence[ |
| 353 | + Union[BatchFeatureView, StreamFeatureView, FeatureView] |
| 354 | + ], |
221 | 355 | entities_to_delete: Sequence[Entity], |
222 | 356 | entities_to_keep: Sequence[Entity], |
223 | 357 | partial: bool, |
@@ -238,20 +372,32 @@ def update( |
238 | 372 | } |
239 | 373 | We remove any feature views named in tables_to_delete. |
240 | 374 | The Entities are serialized in the _id. No schema needs be adjusted. |
| 375 | +
|
| 376 | + When ``vector_enabled`` is set in the online store config, Atlas Vector |
| 377 | + Search indexes are automatically created for feature views containing |
| 378 | + fields with ``vector_index=True`` and dropped for deleted feature views. |
241 | 379 | """ |
242 | 380 | if not isinstance(config.online_store, MongoDBOnlineStoreConfig): |
243 | 381 | raise RuntimeError(f"{config.online_store.type = }. It must be mongodb.") |
244 | 382 |
|
| 383 | + online_config = config.online_store |
245 | 384 | clxn = self._get_collection(repo_config=config) |
246 | 385 |
|
| 386 | + # --- Remove deleted feature views (data + vector search indexes) --- |
247 | 387 | if tables_to_delete: |
248 | 388 | unset_fields = {} |
249 | 389 | for fv in tables_to_delete: |
250 | 390 | unset_fields[f"features.{fv.name}"] = "" |
251 | 391 | unset_fields[f"event_timestamps.{fv.name}"] = "" |
252 | | - |
253 | 392 | clxn.update_many({}, {"$unset": unset_fields}) |
254 | 393 |
|
| 394 | + if online_config.vector_enabled: |
| 395 | + self._drop_vector_indexes_for_tables(clxn, tables_to_delete) |
| 396 | + |
| 397 | + # --- Create vector search indexes for kept feature views --- |
| 398 | + if online_config.vector_enabled: |
| 399 | + self._ensure_vector_indexes(clxn, tables_to_keep, online_config) |
| 400 | + |
255 | 401 | # Note: entities_to_delete contains Entity definitions (metadata), not entity instances. |
256 | 402 | # Like other online stores, we don't need to do anything with entities_to_delete here. |
257 | 403 |
|
@@ -286,6 +432,136 @@ async def close(self) -> None: |
286 | 432 | # Helpers |
287 | 433 | # ------------------------------------------------------------------ |
288 | 434 |
|
| 435 | + # -- Vector Search helpers ------------------------------------------ |
| 436 | + |
| 437 | + @staticmethod |
| 438 | + def _vector_search_index_name(fv_name: str, field_name: str) -> str: |
| 439 | + """Canonical Atlas vector search index name for a (feature_view, field) pair.""" |
| 440 | + return f"{fv_name}__{field_name}__vs_index" |
| 441 | + |
| 442 | + @staticmethod |
| 443 | + def _vector_search_index_definition( |
| 444 | + path: str, |
| 445 | + num_dimensions: int, |
| 446 | + similarity: str, |
| 447 | + ) -> dict: |
| 448 | + """Return a vector search index definition for ``SearchIndexModel``.""" |
| 449 | + return { |
| 450 | + "fields": [ |
| 451 | + { |
| 452 | + "type": "vector", |
| 453 | + "path": path, |
| 454 | + "numDimensions": num_dimensions, |
| 455 | + "similarity": similarity, |
| 456 | + } |
| 457 | + ] |
| 458 | + } |
| 459 | + |
| 460 | + @staticmethod |
| 461 | + def _wait_for_index_ready( |
| 462 | + collection: Collection, |
| 463 | + index_name: str, |
| 464 | + timeout: int, |
| 465 | + poll_interval: float = 1.0, |
| 466 | + ) -> None: |
| 467 | + """Poll until the named Atlas Search index reaches READY status. |
| 468 | +
|
| 469 | + Raises ``TimeoutError`` if the index does not become queryable |
| 470 | + within *timeout* seconds. |
| 471 | + """ |
| 472 | + deadline = time.monotonic() + timeout |
| 473 | + while time.monotonic() < deadline: |
| 474 | + for idx in collection.list_search_indexes(name=index_name): |
| 475 | + if idx.get("status") == "READY" or idx.get("queryable") is True: |
| 476 | + return |
| 477 | + time.sleep(poll_interval) |
| 478 | + raise TimeoutError( |
| 479 | + f"Atlas search index '{index_name}' did not reach READY " |
| 480 | + f"within {timeout}s. Increase vector_index_wait_timeout in " |
| 481 | + f"MongoDBOnlineStoreConfig if the index needs more time." |
| 482 | + ) |
| 483 | + |
| 484 | + def _drop_vector_indexes_for_tables( |
| 485 | + self, |
| 486 | + collection: Collection, |
| 487 | + tables: Sequence[FeatureView], |
| 488 | + ) -> None: |
| 489 | + """Drop all Atlas vector search indexes belonging to the given feature views.""" |
| 490 | + existing = {idx["name"] for idx in collection.list_search_indexes()} |
| 491 | + for fv in tables: |
| 492 | + for field in fv.features: |
| 493 | + idx_name = self._vector_search_index_name(fv.name, field.name) |
| 494 | + if idx_name in existing: |
| 495 | + logger.info("Dropping vector search index: %s", idx_name) |
| 496 | + collection.drop_search_index(idx_name) |
| 497 | + |
| 498 | + def _ensure_vector_indexes( |
| 499 | + self, |
| 500 | + collection: Collection, |
| 501 | + tables: Sequence[Union[BatchFeatureView, StreamFeatureView, FeatureView]], |
| 502 | + online_config: MongoDBOnlineStoreConfig, |
| 503 | + ) -> None: |
| 504 | + """Create Atlas vector search indexes for vector-indexed fields if they don't exist. |
| 505 | +
|
| 506 | + Currently creates one index per (feature_view, vector_field) pair. |
| 507 | + A future optimization could consolidate all vector fields into a |
| 508 | + single composite index with multiple field definitions, reducing |
| 509 | + cluster-wide index count and memory overhead. See Atlas limits: |
| 510 | + hard cap of 2,500 search indexes per cluster; smaller tiers (M10) |
| 511 | + may degrade well before that. |
| 512 | + """ |
| 513 | + db = collection.database |
| 514 | + if collection.name not in db.list_collection_names(): |
| 515 | + db.create_collection(collection.name) |
| 516 | + |
| 517 | + existing = {idx["name"] for idx in collection.list_search_indexes()} |
| 518 | + |
| 519 | + for fv in tables: |
| 520 | + vector_fields = [f for f in fv.features if f.vector_index] |
| 521 | + for field in vector_fields: |
| 522 | + idx_name = self._vector_search_index_name(fv.name, field.name) |
| 523 | + if idx_name in existing: |
| 524 | + logger.debug("Vector search index '%s' already exists", idx_name) |
| 525 | + continue |
| 526 | + |
| 527 | + path = f"features.{fv.name}.{field.name}" |
| 528 | + num_dimensions = field.vector_length |
| 529 | + if not num_dimensions: |
| 530 | + raise ValueError( |
| 531 | + f"Field '{field.name}' in feature view '{fv.name}' has " |
| 532 | + f"vector_index=True but vector_length is not set." |
| 533 | + ) |
| 534 | + |
| 535 | + similarity = ( |
| 536 | + field.vector_search_metric or online_config.similarity or "cosine" |
| 537 | + ) |
| 538 | + |
| 539 | + definition = self._vector_search_index_definition( |
| 540 | + path, num_dimensions, similarity |
| 541 | + ) |
| 542 | + search_index_model = SearchIndexModel( |
| 543 | + definition=definition, |
| 544 | + name=idx_name, |
| 545 | + type="vectorSearch", |
| 546 | + ) |
| 547 | + logger.info( |
| 548 | + "Creating Atlas vector search index '%s' on path '%s' " |
| 549 | + "(dims=%d, similarity=%s)", |
| 550 | + idx_name, |
| 551 | + path, |
| 552 | + num_dimensions, |
| 553 | + similarity, |
| 554 | + ) |
| 555 | + collection.create_search_index(model=search_index_model) |
| 556 | + self._wait_for_index_ready( |
| 557 | + collection, |
| 558 | + idx_name, |
| 559 | + online_config.vector_index_wait_timeout, |
| 560 | + online_config.vector_index_wait_poll_interval, |
| 561 | + ) |
| 562 | + |
| 563 | + # -- Connection helpers --------------------------------------------- |
| 564 | + |
289 | 565 | def _get_client(self, config: RepoConfig): |
290 | 566 | """Returns a connection to the server.""" |
291 | 567 | online_store_config = config.online_store |
@@ -494,5 +770,20 @@ async def online_write_batch_async( |
494 | 770 | progress(len(data)) |
495 | 771 |
|
496 | 772 |
|
497 | | -# TODO |
498 | | -# - Vector Search (requires atlas image in testcontainers or similar) |
| 773 | +def _python_value_to_proto(value: Any) -> ValueProto: |
| 774 | + """Best-effort conversion of a single Python value to ValueProto.""" |
| 775 | + if isinstance(value, float): |
| 776 | + return ValueProto(float_val=value) |
| 777 | + elif isinstance(value, bool): |
| 778 | + return ValueProto(bool_val=value) |
| 779 | + elif isinstance(value, int): |
| 780 | + return ValueProto(int64_val=value) |
| 781 | + elif isinstance(value, str): |
| 782 | + return ValueProto(string_val=value) |
| 783 | + elif isinstance(value, bytes): |
| 784 | + return ValueProto(bytes_val=value) |
| 785 | + elif isinstance(value, list) and all(isinstance(v, float) for v in value): |
| 786 | + proto = ValueProto() |
| 787 | + proto.float_list_val.val.extend(value) |
| 788 | + return proto |
| 789 | + return ValueProto() |