diff --git a/sdk/python/feast/infra/online_stores/elasticsearch_online_store/elasticsearch.py b/sdk/python/feast/infra/online_stores/elasticsearch_online_store/elasticsearch.py index e29ba9df49a..1a0478497d3 100644 --- a/sdk/python/feast/infra/online_stores/elasticsearch_online_store/elasticsearch.py +++ b/sdk/python/feast/infra/online_stores/elasticsearch_online_store/elasticsearch.py @@ -22,6 +22,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 @@ -58,6 +59,13 @@ class ElasticSearchOnlineStoreConfig(FeastConfigBaseModel, VectorStoreConfig): logger = logging.getLogger(__name__) +def _versioned_index_name(table: FeatureView, config: RepoConfig) -> str: + """Return the index name with a version suffix when versioning is enabled.""" + return compute_versioned_name( + table, config.registry.enable_online_feature_view_versioning + ) + + class ElasticsearchFilterTranslator(FilterTranslator): """Translates Feast filters into Elasticsearch Query DSL clauses.""" @@ -179,10 +187,10 @@ def _get_client(self, config: RepoConfig) -> Elasticsearch: ) return self._client - def _bulk_batch_actions(self, table: FeatureView, batch: List[Dict[str, Any]]): + def _bulk_batch_actions(self, index_name: str, batch: List[Dict[str, Any]]): for row in batch: yield { - "_index": table.name, + "_index": index_name, "_id": f"{row['entity_key']}_{row['timestamp']}", "_source": row, } @@ -197,7 +205,8 @@ def online_write_batch( progress: Optional[Callable[[int], Any]], ) -> None: insert_values = [] - include_value_num = self._index_has_value_num(config, table.name) + index_name = _versioned_index_name(table, config) + include_value_num = self._index_has_value_num(config, index_name) grouped_docs: dict[str, dict[str, Any]] = defaultdict( lambda: { "features": {}, @@ -238,7 +247,7 @@ def online_write_batch( batch_size = config.online_store.write_batch_size for i in range(0, len(insert_values), batch_size): batch = insert_values[i : i + batch_size] - actions = self._bulk_batch_actions(table, batch) + actions = self._bulk_batch_actions(index_name, batch) helpers.bulk(self._get_client(config), actions, refresh="wait_for") def online_read( @@ -269,7 +278,9 @@ def online_read( }, } - response = self._get_client(config).search(index=table.name, body=body) + response = self._get_client(config).search( + index=_versioned_index_name(table, config), body=body + ) results = [] @@ -351,7 +362,7 @@ def create_index(self, config: RepoConfig, table: FeatureView): } self._get_client(config).indices.create( - index=table.name, + index=_versioned_index_name(table, config), mappings=index_mapping, ) @@ -366,7 +377,9 @@ def update( ): # implement the update method for table in tables_to_delete: - self._get_client(config).delete_by_query(index=table.name) + self._get_client(config).delete_by_query( + index=_versioned_index_name(table, config) + ) for table in tables_to_keep: self.create_index(config, table) @@ -379,7 +392,9 @@ def teardown( project = config.project try: for table in tables: - self._get_client(config).indices.delete(index=table.name) + self._get_client(config).indices.delete( + index=_versioned_index_name(table, config) + ) except Exception as e: logging.exception(f"Error deleting index in project {project}: {e}") raise @@ -425,7 +440,9 @@ def retrieve_online_documents( } } body = {"size": top_k, "_source": True, "query": query} - response = self._get_client(config).search(index=table.name, body=body) + response = self._get_client(config).search( + index=_versioned_index_name(table, config), body=body + ) rows = response["hits"]["hits"][0:top_k] for row in rows: entity_key = row["_source"]["entity_key"] @@ -485,7 +502,7 @@ def retrieve_online_documents_v2( if embedding is None and query_string is None: raise ValueError("Either embedding or query_string must be provided") - es_index = table.name + es_index = _versioned_index_name(table, config) body: Dict[str, Any] = { "size": top_k, } diff --git a/sdk/python/feast/infra/online_stores/online_store.py b/sdk/python/feast/infra/online_stores/online_store.py index cdf06639fe0..ba919990ee9 100644 --- a/sdk/python/feast/infra/online_stores/online_store.py +++ b/sdk/python/feast/infra/online_stores/online_store.py @@ -315,6 +315,10 @@ def _is_versioned_read_supported(self) -> bool: "feast.infra.online_stores.milvus_online_store.milvus", "MilvusOnlineStore", ), + ( + "feast.infra.online_stores.elasticsearch_online_store.elasticsearch", + "ElasticSearchOnlineStore", + ), ): try: import importlib diff --git a/sdk/python/tests/unit/infra/online_store/test_elasticsearch_versioning.py b/sdk/python/tests/unit/infra/online_store/test_elasticsearch_versioning.py new file mode 100644 index 00000000000..3c1c884f803 --- /dev/null +++ b/sdk/python/tests/unit/infra/online_store/test_elasticsearch_versioning.py @@ -0,0 +1,245 @@ +"""Unit tests for Elasticsearch online store feature view versioning.""" + +from datetime import datetime, timedelta +from unittest.mock import MagicMock, patch + +from feast import Entity, FeatureView +from feast.field import Field +from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + ElasticSearchOnlineStore, +) +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 Float32 +from feast.value_type import ValueType + +MODULE = "feast.infra.online_stores.elasticsearch_online_store.elasticsearch" + + +def _make_feature_view(name="driver_stats", version_number=None, version_tag=None): + entity = Entity( + name="driver_id", + join_keys=["driver_id"], + value_type=ValueType.INT64, + ) + fv = FeatureView( + name=name, + entities=[entity], + ttl=timedelta(days=1), + schema=[Field(name="trips_today", dtype=Float32)], + ) + if version_number is not None: + fv.current_version_number = version_number + if version_tag is not None: + fv.projection.version_tag = version_tag + return fv + + +def _make_config(project="test_project", versioning=False): + config = MagicMock() + config.project = project + config.entity_key_serialization_version = 2 + config.registry.enable_online_feature_view_versioning = versioning + config.online_store.write_batch_size = 100 + config.online_store.similarity = "cosine" + config.online_store.enable_openai_compatible_store = False + return config + + +def _entity_key(): + ek = EntityKeyProto() + ek.join_keys.append("driver_id") + value = ValueProto() + value.int64_val = 1001 + ek.entity_values.append(value) + return ek + + +def _write_one(store, config, fv): + value = ValueProto() + value.float_val = 1.0 + store.online_write_batch( + config, + fv, + [(_entity_key(), {"trips_today": value}, datetime(2026, 1, 1), None)], + None, + ) + + +class TestVersionedIndexName: + """_versioned_index_name names the index the store should touch.""" + + def test_no_versioning(self): + from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + _versioned_index_name, + ) + + fv = _make_feature_view() + config = _make_config(versioning=False) + assert _versioned_index_name(fv, config) == "driver_stats" + + def test_versioning_disabled_ignores_version(self): + from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + _versioned_index_name, + ) + + fv = _make_feature_view(version_number=3) + config = _make_config(versioning=False) + assert _versioned_index_name(fv, config) == "driver_stats" + + def test_versioning_enabled_no_version_set(self): + from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + _versioned_index_name, + ) + + fv = _make_feature_view() + config = _make_config(versioning=True) + assert _versioned_index_name(fv, config) == "driver_stats" + + def test_versioning_enabled_with_current_version_number(self): + from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + _versioned_index_name, + ) + + fv = _make_feature_view(version_number=2) + config = _make_config(versioning=True) + assert _versioned_index_name(fv, config) == "driver_stats_v2" + + def test_version_zero_no_suffix(self): + from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + _versioned_index_name, + ) + + fv = _make_feature_view(version_number=0) + config = _make_config(versioning=True) + assert _versioned_index_name(fv, config) == "driver_stats" + + def test_projection_version_tag_takes_priority(self): + from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + _versioned_index_name, + ) + + fv = _make_feature_view(version_number=1, version_tag=3) + config = _make_config(versioning=True) + assert _versioned_index_name(fv, config) == "driver_stats_v3" + + def test_projection_version_tag_zero_no_suffix(self): + from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + _versioned_index_name, + ) + + fv = _make_feature_view(version_tag=0, version_number=3) + config = _make_config(versioning=True) + assert _versioned_index_name(fv, config) == "driver_stats" + + def test_two_versions_do_not_share_an_index(self): + from feast.infra.online_stores.elasticsearch_online_store.elasticsearch import ( + _versioned_index_name, + ) + + config = _make_config(versioning=True) + v1 = _versioned_index_name(_make_feature_view(version_number=1), config) + v2 = _versioned_index_name(_make_feature_view(version_number=2), config) + assert v1 != v2 + + +class TestStorePathsUseTheVersionedIndex: + """Every path that names an index must name the versioned one.""" + + def test_write_targets_the_versioned_index(self): + store = ElasticSearchOnlineStore() + config = _make_config(versioning=True) + fv = _make_feature_view(version_number=2) + with ( + patch.object(ElasticSearchOnlineStore, "_get_client", MagicMock()), + patch.object( + ElasticSearchOnlineStore, "_index_has_value_num", return_value=False + ), + patch(f"{MODULE}.helpers") as es_helpers, + ): + _write_one(store, config, fv) + actions = list(es_helpers.bulk.call_args[0][1]) + assert actions + assert {action["_index"] for action in actions} == {"driver_stats_v2"} + + def test_write_is_unchanged_when_versioning_is_off(self): + store = ElasticSearchOnlineStore() + config = _make_config(versioning=False) + fv = _make_feature_view(version_number=2) + with ( + patch.object(ElasticSearchOnlineStore, "_get_client", MagicMock()), + patch.object( + ElasticSearchOnlineStore, "_index_has_value_num", return_value=False + ), + patch(f"{MODULE}.helpers") as es_helpers, + ): + _write_one(store, config, fv) + actions = list(es_helpers.bulk.call_args[0][1]) + assert {action["_index"] for action in actions} == {"driver_stats"} + + def test_write_checks_the_mapping_of_the_versioned_index(self): + store = ElasticSearchOnlineStore() + config = _make_config(versioning=True) + fv = _make_feature_view(version_number=2) + with ( + patch.object(ElasticSearchOnlineStore, "_get_client", MagicMock()), + patch.object( + ElasticSearchOnlineStore, "_index_has_value_num", return_value=False + ) as has_value_num, + patch(f"{MODULE}.helpers"), + ): + _write_one(store, config, fv) + assert has_value_num.call_args[0][1] == "driver_stats_v2" + + def test_read_searches_the_versioned_index(self): + store = ElasticSearchOnlineStore() + config = _make_config(versioning=True) + fv = _make_feature_view(version_number=2) + client = MagicMock() + client.search.return_value = {"hits": {"hits": []}} + with patch.object(ElasticSearchOnlineStore, "_get_client", return_value=client): + store.online_read(config, fv, [_entity_key()], ["trips_today"]) + assert client.search.call_args.kwargs["index"] == "driver_stats_v2" + + def test_create_index_creates_the_versioned_index(self): + store = ElasticSearchOnlineStore() + config = _make_config(versioning=True) + fv = _make_feature_view(version_number=2) + client = MagicMock() + with patch.object(ElasticSearchOnlineStore, "_get_client", return_value=client): + store.create_index(config, fv) + assert client.indices.create.call_args.kwargs["index"] == "driver_stats_v2" + + def test_update_deletes_and_creates_versioned_indices(self): + store = ElasticSearchOnlineStore() + config = _make_config(versioning=True) + keep = _make_feature_view(name="keep_me", version_number=2) + drop = _make_feature_view(name="drop_me", version_number=3) + client = MagicMock() + with patch.object(ElasticSearchOnlineStore, "_get_client", return_value=client): + store.update(config, [drop], [keep], [], [], partial=False) + assert client.delete_by_query.call_args.kwargs["index"] == "drop_me_v3" + assert client.indices.create.call_args.kwargs["index"] == "keep_me_v2" + + def test_teardown_deletes_the_versioned_index(self): + store = ElasticSearchOnlineStore() + config = _make_config(versioning=True) + fv = _make_feature_view(version_number=2) + client = MagicMock() + with patch.object(ElasticSearchOnlineStore, "_get_client", return_value=client): + store.teardown(config, [fv], []) + assert client.indices.delete.call_args.kwargs["index"] == "driver_stats_v2" + + +class TestElasticsearchVersionedReadSupport: + """The store must also be on the base class's supported list: without that, + every versioned read is refused however the indices are named.""" + + def test_store_declares_versioned_read_support(self): + assert ElasticSearchOnlineStore()._is_versioned_read_supported() is True + + def test_versioned_ref_is_not_refused(self): + store = ElasticSearchOnlineStore() + fv = _make_feature_view() + fv.projection.version_tag = 2 + store._check_versioned_read_support([(fv, ["trips_today"])])