From 0800f135332adf57b68a7fb95ac2faee9125f77f Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sun, 14 Jun 2026 15:38:45 +0530 Subject: [PATCH 01/16] feat: Implement Databricks Unity Catalog offline store integration Signed-off-by: Abhishek Shinde --- .../spark_offline_store/databricks_uc.py | 300 ++++++++++++++++++ sdk/python/feast/repo_config.py | 1 + .../spark_offline_store/test_databricks_uc.py | 193 +++++++++++ 3 files changed, 494 insertions(+) create mode 100644 sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py create mode 100644 sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py new file mode 100644 index 00000000000..e6eb3ea47e6 --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py @@ -0,0 +1,300 @@ +import logging +from datetime import date, datetime +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import pandas as pd +import pyarrow +import pyspark +from pydantic import StrictStr +from pyspark import SparkConf +from pyspark.sql import SparkSession + +from feast import FeatureView +from feast.data_source import DataSource +from feast.infra.offline_stores.contrib.spark_offline_store.spark import ( + SparkOfflineStore, + SparkOfflineStoreConfig, +) +from feast.infra.offline_stores.offline_store import RetrievalJob +from feast.infra.registry.base_registry import BaseRegistry +from feast.repo_config import RepoConfig + +logger = logging.getLogger(__name__) + + +class DatabricksUCOfflineStoreConfig(SparkOfflineStoreConfig): + type: StrictStr = "databricks_uc" + """Offline store type selector""" + + workspace_host: Optional[StrictStr] = None + """Databricks workspace host (e.g. adb-xxxx.azuredatabricks.net)""" + + token: Optional[StrictStr] = None + """Databricks Personal Access Token (PAT)""" + + cluster_id: Optional[StrictStr] = None + """Databricks Cluster ID to connect to for Databricks Connect""" + + default_catalog: Optional[StrictStr] = None + """Default catalog name to use in Unity Catalog""" + + default_schema: Optional[StrictStr] = None + """Default schema name to use in Unity Catalog""" + + +def get_databricks_session( + store_config: DatabricksUCOfflineStoreConfig, +) -> SparkSession: + # Check if there is already an active session + spark_session = SparkSession.getActiveSession() + if not spark_session: + workspace_host = store_config.workspace_host + token = store_config.token + cluster_id = store_config.cluster_id + + # Clean host URL if it starts with https:// + if workspace_host: + if workspace_host.startswith("https://"): + workspace_host = workspace_host[8:] + elif workspace_host.startswith("http://"): + workspace_host = workspace_host[7:] + + if workspace_host and cluster_id: + # Databricks Connect V2 initialization (Spark Connect URI format) + conn_str = f"sc://{workspace_host}:443/" + params = [] + if token: + params.append(f"token={token}") + params.append(f"x-databricks-cluster-id={cluster_id}") + if params: + conn_str = f"{conn_str};{';'.join(params)}" + + try: + from databricks.connect import DatabricksSession + + builder = DatabricksSession.builder.remote(conn_str) + except ImportError: + # Fallback to standard PySpark remote connect if databricks-connect not installed + builder = SparkSession.builder.remote(conn_str) + else: + try: + from databricks.connect import DatabricksSession + + builder = DatabricksSession.builder + except ImportError: + builder = SparkSession.builder + + spark_conf = store_config.spark_conf + if spark_conf: + builder = builder.config( + conf=SparkConf().setAll([(k, v) for k, v in spark_conf.items()]) + ) + + spark_session = builder.getOrCreate() + + # Apply configuration defaults + spark_session.conf.set("spark.sql.parser.quotedRegexColumnNames", "true") + + if store_config.default_catalog: + spark_session.sql(f"USE CATALOG `{store_config.default_catalog}`") + if store_config.default_schema: + spark_session.sql(f"USE SCHEMA `{store_config.default_schema}`") + + return spark_session + + +class DatabricksUCOfflineStore(SparkOfflineStore): + @staticmethod + def pull_latest_from_table_or_query( + config: RepoConfig, + data_source: DataSource, + join_key_columns: List[str], + feature_name_columns: List[str], + timestamp_field: str, + created_timestamp_column: Optional[str], + start_date: datetime, + end_date: datetime, + ) -> RetrievalJob: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + # Initialize/Retrieve the Databricks Spark Session so it's registered as active + get_databricks_session(config.offline_store) + + return SparkOfflineStore.pull_latest_from_table_or_query( + config=config, + data_source=data_source, + join_key_columns=join_key_columns, + feature_name_columns=feature_name_columns, + timestamp_field=timestamp_field, + created_timestamp_column=created_timestamp_column, + start_date=start_date, + end_date=end_date, + ) + + @staticmethod + def get_historical_features( + config: RepoConfig, + feature_views: List[FeatureView], + feature_refs: List[str], + entity_df: Optional[Union[pd.DataFrame, str, pyspark.sql.DataFrame]], + registry: BaseRegistry, + project: str, + full_feature_names: bool = False, + **kwargs, + ) -> RetrievalJob: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.get_historical_features( + config=config, + feature_views=feature_views, + feature_refs=feature_refs, + entity_df=entity_df, + registry=registry, + project=project, + full_feature_names=full_feature_names, + **kwargs, + ) + + @staticmethod + def pull_all_from_table_or_query( + config: RepoConfig, + data_source: DataSource, + join_key_columns: List[str], + feature_name_columns: List[str], + timestamp_field: str, + created_timestamp_column: Optional[str] = None, + start_date: Optional[datetime] = None, + end_date: Optional[datetime] = None, + ) -> RetrievalJob: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.pull_all_from_table_or_query( + config=config, + data_source=data_source, + join_key_columns=join_key_columns, + feature_name_columns=feature_name_columns, + timestamp_field=timestamp_field, + created_timestamp_column=created_timestamp_column, + start_date=start_date, + end_date=end_date, + ) + + @staticmethod + def offline_write_batch( + config: RepoConfig, + feature_view: FeatureView, + table: pyarrow.Table, + progress: Optional[Callable[[int], Any]], + ): + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.offline_write_batch( + config=config, + feature_view=feature_view, + table=table, + progress=progress, + ) + + @staticmethod + def compute_monitoring_metrics( + config: RepoConfig, + data_source: DataSource, + feature_columns: List[Tuple[str, str]], + timestamp_field: str, + start_date: Optional[datetime] = None, + end_date: Optional[datetime] = None, + histogram_bins: int = 20, + top_n: int = 10, + ) -> List[Dict[str, Any]]: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.compute_monitoring_metrics( + config=config, + data_source=data_source, + feature_columns=feature_columns, + timestamp_field=timestamp_field, + start_date=start_date, + end_date=end_date, + histogram_bins=histogram_bins, + top_n=top_n, + ) + + @staticmethod + def get_monitoring_max_timestamp( + config: RepoConfig, + data_source: DataSource, + timestamp_field: str, + ) -> Optional[datetime]: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.get_monitoring_max_timestamp( + config=config, + data_source=data_source, + timestamp_field=timestamp_field, + ) + + @staticmethod + def ensure_monitoring_tables(config: RepoConfig) -> None: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.ensure_monitoring_tables(config=config) + + @staticmethod + def save_monitoring_metrics( + config: RepoConfig, + metric_type: str, + metrics: List[Dict[str, Any]], + ) -> None: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.save_monitoring_metrics( + config=config, + metric_type=metric_type, + metrics=metrics, + ) + + @staticmethod + def query_monitoring_metrics( + config: RepoConfig, + project: str, + metric_type: str, + filters: Optional[Dict[str, Any]] = None, + start_date: Optional[date] = None, + end_date: Optional[date] = None, + ) -> List[Dict[str, Any]]: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.query_monitoring_metrics( + config=config, + project=project, + metric_type=metric_type, + filters=filters, + start_date=start_date, + end_date=end_date, + ) + + @staticmethod + def clear_monitoring_baseline( + config: RepoConfig, + project: str, + feature_view_name: Optional[str] = None, + feature_name: Optional[str] = None, + data_source_type: Optional[str] = None, + ) -> None: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + + return SparkOfflineStore.clear_monitoring_baseline( + config=config, + project=project, + feature_view_name=feature_view_name, + feature_name=feature_name, + data_source_type=data_source_type, + ) diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index 06529cea0f2..8527a3591aa 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -95,6 +95,7 @@ "redshift": "feast.infra.offline_stores.redshift.RedshiftOfflineStore", "snowflake.offline": "feast.infra.offline_stores.snowflake.SnowflakeOfflineStore", "spark": "feast.infra.offline_stores.contrib.spark_offline_store.spark.SparkOfflineStore", + "databricks_uc": "feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc.DatabricksUCOfflineStore", "trino": "feast.infra.offline_stores.contrib.trino_offline_store.trino.TrinoOfflineStore", "postgres": "feast.infra.offline_stores.contrib.postgres_offline_store.postgres.PostgreSQLOfflineStore", "athena": "feast.infra.offline_stores.contrib.athena_offline_store.athena.AthenaOfflineStore", diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py new file mode 100644 index 00000000000..07bfd1d5416 --- /dev/null +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py @@ -0,0 +1,193 @@ +from datetime import datetime, timezone +from unittest.mock import MagicMock, patch + +import pytest +from pydantic import ValidationError + +from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( + DatabricksUCOfflineStore, + DatabricksUCOfflineStoreConfig, + get_databricks_session, +) +from feast.infra.offline_stores.contrib.spark_offline_store.spark_source import ( + SparkSource, +) +from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig +from feast.repo_config import RepoConfig + + +def test_config_parsing(): + config_dict = { + "type": "databricks_uc", + "workspace_host": "adb-12345.azuredatabricks.net", + "token": "dapi123456", + "cluster_id": "0123-4567-abcde", + "default_catalog": "main", + "default_schema": "default", + "spark_conf": {"spark.sql.shuffle.partitions": "10"}, + } + config = DatabricksUCOfflineStoreConfig(**config_dict) + assert config.type == "databricks_uc" + assert config.workspace_host == "adb-12345.azuredatabricks.net" + assert config.token == "dapi123456" + assert config.cluster_id == "0123-4567-abcde" + assert config.default_catalog == "main" + assert config.default_schema == "default" + assert config.spark_conf == {"spark.sql.shuffle.partitions": "10"} + + +def test_config_forbidden_extra(): + with pytest.raises(ValidationError): + DatabricksUCOfflineStoreConfig(type="databricks_uc", invalid_key="some_val") + + +@patch("pyspark.sql.SparkSession.getActiveSession") +def test_get_databricks_session_active(mock_get_active): + mock_session = MagicMock() + mock_get_active.return_value = mock_session + + config = DatabricksUCOfflineStoreConfig( + type="databricks_uc", + default_catalog="my_catalog", + default_schema="my_schema", + ) + + session = get_databricks_session(config) + + assert session == mock_session + mock_session.conf.set.assert_called_once_with( + "spark.sql.parser.quotedRegexColumnNames", "true" + ) + mock_session.sql.assert_any_call("USE CATALOG `my_catalog`") + mock_session.sql.assert_any_call("USE SCHEMA `my_schema`") + + +@patch("pyspark.sql.SparkSession.getActiveSession") +@patch("pyspark.sql.SparkSession.builder") +def test_get_databricks_session_new_remote(mock_builder, mock_get_active): + mock_get_active.return_value = None + mock_session = MagicMock() + mock_builder.remote.return_value.config.return_value.getOrCreate.return_value = ( + mock_session + ) + + config = DatabricksUCOfflineStoreConfig( + type="databricks_uc", + workspace_host="https://adb-12345.azuredatabricks.net", + token="dapi123", + cluster_id="0123-4567-abcde", + spark_conf={"spark.some.option": "value"}, + ) + + session = get_databricks_session(config) + + assert session == mock_session + mock_builder.remote.assert_called_once_with( + "sc://adb-12345.azuredatabricks.net:443/;token=dapi123;x-databricks-cluster-id=0123-4567-abcde" + ) + + +@patch( + "feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc.get_databricks_session" +) +@patch( + "feast.infra.offline_stores.contrib.spark_offline_store.spark.SparkOfflineStore.get_historical_features" +) +def test_get_historical_features_delegation(mock_parent_features, mock_get_session): + mock_session = MagicMock() + mock_get_session.return_value = mock_session + + repo_config = RepoConfig( + registry="file:///tmp/registry.db", + project="test", + provider="local", + online_store=SqliteOnlineStoreConfig(type="sqlite"), + offline_store=DatabricksUCOfflineStoreConfig( + type="databricks_uc", + workspace_host="adb-123.databricks.com", + cluster_id="123", + ), + ) + + feature_views = [] + feature_refs = ["fv:f1"] + entity_df = MagicMock() + registry = MagicMock() + + DatabricksUCOfflineStore.get_historical_features( + config=repo_config, + feature_views=feature_views, + feature_refs=feature_refs, + entity_df=entity_df, + registry=registry, + project="test", + ) + + mock_get_session.assert_called_once_with(repo_config.offline_store) + mock_parent_features.assert_called_once_with( + config=repo_config, + feature_views=feature_views, + feature_refs=feature_refs, + entity_df=entity_df, + registry=registry, + project="test", + full_feature_names=False, + ) + + +@patch( + "feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc.get_databricks_session" +) +@patch( + "feast.infra.offline_stores.contrib.spark_offline_store.spark.SparkOfflineStore.pull_latest_from_table_or_query" +) +def test_pull_latest_from_table_or_query_delegation( + mock_parent_pull_latest, mock_get_session +): + mock_session = MagicMock() + mock_get_session.return_value = mock_session + + repo_config = RepoConfig( + registry="file:///tmp/registry.db", + project="test", + provider="local", + online_store=SqliteOnlineStoreConfig(type="sqlite"), + offline_store=DatabricksUCOfflineStoreConfig( + type="databricks_uc", + workspace_host="adb-123.databricks.com", + cluster_id="123", + ), + ) + + data_source = SparkSource( + name="test_source", + path="catalog.schema.table", + file_format="parquet", + timestamp_field="ts", + ) + + start_date = datetime(2023, 1, 1, tzinfo=timezone.utc) + end_date = datetime(2023, 1, 2, tzinfo=timezone.utc) + + DatabricksUCOfflineStore.pull_latest_from_table_or_query( + config=repo_config, + data_source=data_source, + join_key_columns=["id"], + feature_name_columns=["val"], + timestamp_field="ts", + created_timestamp_column=None, + start_date=start_date, + end_date=end_date, + ) + + mock_get_session.assert_called_once_with(repo_config.offline_store) + mock_parent_pull_latest.assert_called_once_with( + config=repo_config, + data_source=data_source, + join_key_columns=["id"], + feature_name_columns=["val"], + timestamp_field="ts", + created_timestamp_column=None, + start_date=start_date, + end_date=end_date, + ) From 664107f0a27da3c2adf0aadeeacd9059c2a014d4 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Tue, 16 Jun 2026 09:27:15 +0530 Subject: [PATCH 02/16] fix: Initialize Databricks session in DatabricksUCOfflineStore validation methods Signed-off-by: Abhishek Shinde --- .../spark_offline_store/databricks_uc.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py index e6eb3ea47e6..9c8ec25e5d3 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py @@ -1,6 +1,6 @@ import logging from datetime import date, datetime -from typing import Any, Callable, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union import pandas as pd import pyarrow @@ -298,3 +298,21 @@ def clear_monitoring_baseline( feature_name=feature_name, data_source_type=data_source_type, ) + + @staticmethod + def validate_data_source( + config: RepoConfig, + data_source: DataSource, + ): + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + data_source.validate(config=config) + + @staticmethod + def get_table_column_names_and_types_from_data_source( + config: RepoConfig, + data_source: DataSource, + ) -> Iterable[Tuple[str, str]]: + assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) + get_databricks_session(config.offline_store) + return data_source.get_table_column_names_and_types(config=config) From 3b30044732ae09ae61978002e8d8baf4939fed83 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sun, 21 Jun 2026 10:37:34 +0530 Subject: [PATCH 03/16] fix: Fix MyPy errors by asserting SparkSession is not None Signed-off-by: Abhishek Shinde --- .../offline_stores/contrib/spark_offline_store/databricks_uc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py index 9c8ec25e5d3..5c8228710dc 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py @@ -92,6 +92,8 @@ def get_databricks_session( spark_session = builder.getOrCreate() + assert spark_session is not None + # Apply configuration defaults spark_session.conf.set("spark.sql.parser.quotedRegexColumnNames", "true") From faa0675ba46398113bc6f53956a09581719e70c5 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sat, 27 Jun 2026 18:21:22 +0530 Subject: [PATCH 04/16] feat: Register FeatureViews as Unity Catalog feature tables on feast apply Signed-off-by: Abhishek Shinde --- sdk/python/feast/feature_store.py | 70 ++++ .../spark_offline_store/databricks_uc.py | 20 +- .../spark_offline_store/uc_registration.py | 387 ++++++++++++++++++ 3 files changed, 475 insertions(+), 2 deletions(-) create mode 100644 sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 3dcd4189d3c..cdd5af112ba 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -1287,6 +1287,10 @@ def _apply_diffs( # Emit OpenLineage events for applied objects self._emit_openlineage_apply_diffs(registry_diff) + # Register feature views as Unity Catalog feature tables (if using + # the databricks_uc offline store) + self._register_uc_feature_tables_from_diffs(registry_diff) + # Emit MLflow events for applied objects (Phase 7) self._mlflow_log_apply_diffs(registry_diff) @@ -1637,6 +1641,10 @@ def apply( # Emit OpenLineage events for applied objects self._emit_openlineage_apply(objects) + # Register feature views as Unity Catalog feature tables (if using + # the databricks_uc offline store) + self._register_uc_feature_tables_legacy(objects) + # Emit MLflow events for applied objects (Phase 7) self._mlflow_log_apply(objects) @@ -1665,6 +1673,68 @@ def _emit_openlineage_apply(self, objects: List[Any]): except Exception as e: warnings.warn(f"Failed to emit OpenLineage apply events: {e}") + # ------------------------------------------------------------------ # + # Unity Catalog feature table registration hooks + # ------------------------------------------------------------------ # + + def _register_uc_feature_tables_from_diffs( + self, registry_diff: RegistryDiff + ) -> None: + """Register applied feature views as UC feature tables. + + Only active when the offline store is ``databricks_uc`` and + ``uc_registration.enabled`` is True. + """ + from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( + DatabricksUCOfflineStoreConfig, + ) + from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( + register_uc_feature_tables, + ) + + if not isinstance(self.config.offline_store, DatabricksUCOfflineStoreConfig): + return + + uc_config = self.config.offline_store.uc_registration + if uc_config is None or not uc_config.enabled: + return + + fvs = [ + d.new_feast_object + for d in registry_diff.feast_object_diffs + if d.new_feast_object and isinstance(d.new_feast_object, FeatureView) + ] + if not fvs: + return + + register_uc_feature_tables(self.config.offline_store, fvs, self.project) + + def _register_uc_feature_tables_legacy(self, objects: List[Any]) -> None: + """Register feature views as UC feature tables (legacy apply path). + + Only active when the offline store is ``databricks_uc`` and + ``uc_registration.enabled`` is True. + """ + from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( + DatabricksUCOfflineStoreConfig, + ) + from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( + register_uc_feature_tables, + ) + + if not isinstance(self.config.offline_store, DatabricksUCOfflineStoreConfig): + return + + uc_config = self.config.offline_store.uc_registration + if uc_config is None or not uc_config.enabled: + return + + fvs = [obj for obj in objects if isinstance(obj, FeatureView)] + if not fvs: + return + + register_uc_feature_tables(self.config.offline_store, fvs, self.project) + def teardown(self): """Tears down all local and cloud resources for the feature store.""" tables: List[BaseFeatureView] = [] diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py index 5c8228710dc..d90caefa589 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py @@ -5,7 +5,7 @@ import pandas as pd import pyarrow import pyspark -from pydantic import StrictStr +from pydantic import StrictBool, StrictStr from pyspark import SparkConf from pyspark.sql import SparkSession @@ -17,11 +17,24 @@ ) from feast.infra.offline_stores.offline_store import RetrievalJob from feast.infra.registry.base_registry import BaseRegistry -from feast.repo_config import RepoConfig +from feast.repo_config import FeastConfigBaseModel, RepoConfig logger = logging.getLogger(__name__) +class UCRegistrationConfig(FeastConfigBaseModel): + """Configuration for Unity Catalog feature table registration during ``feast apply``.""" + + enabled: StrictBool = True + """Whether to register feature views as UC feature tables on ``feast apply``.""" + + catalog: Optional[StrictStr] = None + """Default catalog for UC feature tables. Overrides ``DatabricksUCOfflineStoreConfig.default_catalog``.""" + + schema: Optional[StrictStr] = None + """Default schema for UC feature tables. Overrides ``DatabricksUCOfflineStoreConfig.default_schema``.""" + + class DatabricksUCOfflineStoreConfig(SparkOfflineStoreConfig): type: StrictStr = "databricks_uc" """Offline store type selector""" @@ -41,6 +54,9 @@ class DatabricksUCOfflineStoreConfig(SparkOfflineStoreConfig): default_schema: Optional[StrictStr] = None """Default schema name to use in Unity Catalog""" + uc_registration: Optional[UCRegistrationConfig] = None + """Configuration for UC feature table registration during ``feast apply``.""" + def get_databricks_session( store_config: DatabricksUCOfflineStoreConfig, diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py new file mode 100644 index 00000000000..87ee8b1a2b9 --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py @@ -0,0 +1,387 @@ +"""Unity Catalog feature table registration for ``feast apply``. + +When the offline store is configured as ``databricks_uc``, this module +registers (or updates) each FeatureView as a Unity Catalog feature table +via the Databricks ``FeatureEngineeringClient``. FeatureView entities +become UC primary keys, and tags/description/owner are synced to UC metadata. + +Per‑FeatureView opt‑out and overrides are controlled via FeatureView ``tags``: + +* ``uc.register_as_feature_table`` — ``"false"`` skips a specific view. +* ``uc.catalog`` / ``uc.schema`` / ``uc.table`` — override the UC path. + +Global defaults come from ``UCRegistrationConfig`` (inside +``DatabricksUCOfflineStoreConfig``). +""" + +import logging +from typing import Dict, List, Optional, Tuple + +import click +from pyspark.sql import SparkSession +from pyspark.sql.types import ( + StructField, + StructType, +) + +from feast import FeatureView +from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( + DatabricksUCOfflineStoreConfig, + get_databricks_session, +) + +logger = logging.getLogger(__name__) + +# FeatureView tag keys for per‑FV UC configuration +_REGISTER_AS_FEATURE_TABLE_KEY = "uc.register_as_feature_table" +_CATALOG_KEY = "uc.catalog" +_SCHEMA_KEY = "uc.schema" +_TABLE_KEY = "uc.table" + +# Internal Feast tags stored on the UC table +_MANAGED_BY_TAG = "feast_managed" + + +def _feast_to_spark_type_simple(field): + """Convert a Feast :class:`Field` to a pyspark :class:`DataType`. + + This is a best‑effort mapping; the underlying Spark‑read path performs + full schema inference at query time. + """ + from pyspark.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + ByteType, + DateType, + DoubleType, + FloatType, + IntegerType, + LongType, + ShortType, + StringType, + TimestampType, + ) + + dtype = getattr(field, "dtype", None) + if dtype is None: + return StringType() # fallback + + type_name = str(dtype).upper() + + type_map = { + "BOOL": BooleanType(), + "BOOLEAN": BooleanType(), + "INT8": ByteType(), + "BYTE": ByteType(), + "INT16": ShortType(), + "SHORT": ShortType(), + "INT32": IntegerType(), + "INT": IntegerType(), + "INTEGER": IntegerType(), + "INT64": LongType(), + "LONG": LongType(), + "BIGINT": LongType(), + "FLOAT32": FloatType(), + "FLOAT": FloatType(), + "FLOAT64": DoubleType(), + "DOUBLE": DoubleType(), + "STRING": StringType(), + "UTF8": StringType(), + "BINARY": BinaryType(), + "TIMESTAMP": TimestampType(), + "TIMESTAMP_TZ": TimestampType(), + "DATE": DateType(), + "DATE32": DateType(), + } + + if type_name.startswith("LIST<") or type_name.startswith("ARRAY<"): + return ArrayType(StringType()) + if "DOUBLE" in type_name or "FLOAT" in type_name and "64" in type_name: + return DoubleType() + if "LIST" in type_name or "ARRAY" in type_name: + return ArrayType(StringType()) + + return type_map.get(type_name, StringType()) + + +def _should_register(fv: FeatureView) -> bool: + """Return ``True`` unless the feature view opts out via tags.""" + return fv.tags.get(_REGISTER_AS_FEATURE_TABLE_KEY, "true").lower() != "false" + + +def _resolve_uc_path( + fv: FeatureView, + default_catalog: Optional[str], + default_schema: Optional[str], +) -> Tuple[str, str, str]: + """Resolve the (catalog, schema, table_name) for a feature view. + + Prioritises per‑FV tag overrides, then global defaults, then the + feature view name as the table name. + """ + catalog = fv.tags.get(_CATALOG_KEY) or default_catalog + schema = fv.tags.get(_SCHEMA_KEY) or default_schema + table = fv.tags.get(_TABLE_KEY) or fv.name + + # Sanitise: replace characters that are invalid in UC names + table = table.replace("-", "_").replace(".", "_").replace(" ", "_") + return catalog, schema, table + + +def _build_spark_schema(fv: FeatureView) -> StructType: + """Build a pyspark ``StructType`` from the feature view's columns.""" + fields: List[StructField] = [] + + seen: set = set() + for col in fv.entity_columns: + if col.name not in seen: + fields.append( + StructField(col.name, _feast_to_spark_type_simple(col), nullable=False) + ) + seen.add(col.name) + + for col in fv.features: + if col.name not in seen: + fields.append( + StructField(col.name, _feast_to_spark_type_simple(col), nullable=True) + ) + seen.add(col.name) + + # Add timestamp column if the batch source declares one + timestamp_field = getattr(fv.batch_source, "timestamp_field", None) + if timestamp_field and timestamp_field not in seen: + from pyspark.sql.types import TimestampType + + fields.append(StructField(timestamp_field, TimestampType(), nullable=True)) + + return StructType(fields) + + +def _get_primary_keys(fv: FeatureView) -> List[str]: + """Extract primary key column names from entity columns.""" + return [col.name for col in fv.entity_columns] + + +def _build_uc_tags(fv: FeatureView, project: str) -> Dict[str, str]: + """Build UC table tags, excluding internal ``uc.*`` keys.""" + tags: Dict[str, str] = {} + for key, value in fv.tags.items(): + if not key.startswith("uc."): + tags[key] = value + tags[_MANAGED_BY_TAG] = "feast" + tags["feast_project"] = project + return tags + + +def _escape_sql_string(value: str) -> str: + """Escape single quotes for SQL string literals.""" + return value.replace("\\", "\\\\").replace("'", "\\'") + + +def _build_full_table_name(catalog, schema, table): + """Build a three-level table reference: ``catalog.schema.table``.""" + parts = [] + if catalog: + parts.append(f"`{catalog}`") + if schema: + parts.append(f"`{schema}`") + parts.append(f"`{table}`") + return ".".join(parts) + + +def _table_exists(spark_session: SparkSession, full_name: str) -> bool: + """Check whether a UC table exists.""" + try: + return spark_session.catalog.tableExists(full_name) + except Exception: + return False + + +def _create_uc_feature_table( + fe_client, + spark_session: SparkSession, + fv: FeatureView, + full_name: str, + primary_keys: List[str], + project: str, +) -> None: + """Create a new UC feature table via FeatureEngineeringClient.""" + spark_schema = _build_spark_schema(fv) + timestamp_key = getattr(fv.batch_source, "timestamp_field", None) + description = fv.description or "" + tags = _build_uc_tags(fv, project) + owner = fv.owner + + fe_client.create_table( + name=full_name, + primary_keys=primary_keys, + schema=spark_schema, + timestamp_key=timestamp_key, + description=description, + tags=tags, + ) + + if owner: + try: + spark_session.sql( + f"ALTER TABLE {full_name} SET OWNER TO `{_escape_sql_string(owner)}`" + ) + except Exception: + logger.debug( + "Could not set owner for UC table %s; continuing", + full_name, + ) + + +def _update_uc_feature_table_metadata( + spark_session: SparkSession, + fv: FeatureView, + full_name: str, + project: str, +) -> None: + """Update metadata on an existing UC feature table.""" + description = fv.description or "" + owner = fv.owner + tags = _build_uc_tags(fv, project) + + # Update description (COMMENT) + try: + spark_session.sql( + f"COMMENT ON TABLE {full_name} IS '{_escape_sql_string(description)}'" + ) + except Exception: + logger.debug( + "Could not update comment on UC table %s; continuing", + full_name, + ) + + # Update tags + if tags: + tag_pairs = ", ".join( + f"'{_escape_sql_string(k)}' = '{_escape_sql_string(v)}'" + for k, v in tags.items() + ) + try: + spark_session.sql(f"ALTER TABLE {full_name} SET TAGS ({tag_pairs})") + except Exception: + logger.debug( + "Could not update tags on UC table %s; continuing", + full_name, + ) + + # Update owner + if owner: + try: + spark_session.sql( + f"ALTER TABLE {full_name} SET OWNER TO `{_escape_sql_string(owner)}`" + ) + except Exception: + logger.debug( + "Could not update owner on UC table %s; continuing", + full_name, + ) + + +def _register_single_feature_view( + spark_session: SparkSession, + fe_client, + fv: FeatureView, + project: str, + default_catalog: Optional[str], + default_schema: Optional[str], +) -> None: + """Register or update a single FeatureView as a UC feature table.""" + if not _should_register(fv): + click.echo(f" ⊘ Skipping UC registration for {fv.name} (opt‑out)") + return + + catalog, schema, table = _resolve_uc_path(fv, default_catalog, default_schema) + if not catalog or not schema: + click.echo( + f" ⚠ Cannot register {fv.name}: missing catalog or schema. " + f"Set default_catalog/default_schema in feature_store.yaml or " + f"use tags uc.catalog/uc.schema." + ) + return + + full_name = _build_full_table_name(catalog, schema, table) + primary_keys = _get_primary_keys(fv) + + if _table_exists(spark_session, full_name): + _update_uc_feature_table_metadata(spark_session, fv, full_name, project) + click.echo(f" ✓ Updated UC feature table: {full_name}") + else: + _create_uc_feature_table( + fe_client, + spark_session, + fv, + full_name, + primary_keys, + project, + ) + click.echo(f" ✓ Created UC feature table: {full_name}") + + +def register_uc_feature_tables( + config: DatabricksUCOfflineStoreConfig, + feature_views: List[FeatureView], + project: str, +) -> None: + """Register or update FeatureViews as Unity Catalog feature tables. + + Skips silently when: + - ``uc_registration`` is not configured or ``enabled`` is ``False``. + - ``databricks-feature-engineering`` is not installed. + - The Databricks Spark session cannot be created. + + Per‑feature‑view errors are logged and do not halt the apply. + """ + uc_config = config.uc_registration + if uc_config is None or not uc_config.enabled: + return + + try: + from databricks.feature_engineering import ( + FeatureEngineeringClient, # noqa: F401 + ) + except ImportError: + logger.info( + "databricks-feature-engineering is not installed; " + "skipping UC feature table registration. " + "Install with: pip install databricks-feature-engineering" + ) + return + + try: + spark_session = get_databricks_session(config) + except Exception as e: + logger.warning( + "Could not create Databricks Spark session for UC registration: %s", e + ) + return + + fe_client = FeatureEngineeringClient() + + default_catalog = uc_config.catalog or config.default_catalog + default_schema = uc_config.schema or config.default_schema + + for fv in feature_views: + try: + _register_single_feature_view( + spark_session, + fe_client, + fv, + project, + default_catalog, + default_schema, + ) + except Exception: + logger.exception( + "Failed to register UC feature table for FeatureView '%s'", + fv.name, + ) + click.echo( + f" ✗ Failed to register UC feature table: {fv.name} " + f"(check logs for details)" + ) From 327228c6f69d8f11bcdf120bdd5836f7e5a18d07 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sun, 28 Jun 2026 15:04:00 +0530 Subject: [PATCH 05/16] fix: Rename schema to uc_schema, add missing type mappings, add tests, fix operator precedence Signed-off-by: Abhishek Shinde --- .../spark_offline_store/databricks_uc.py | 2 +- .../spark_offline_store/uc_registration.py | 8 +- .../test_uc_registration.py | 448 ++++++++++++++++++ 3 files changed, 454 insertions(+), 4 deletions(-) create mode 100644 sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py index d90caefa589..3fd48b67d96 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py @@ -31,7 +31,7 @@ class UCRegistrationConfig(FeastConfigBaseModel): catalog: Optional[StrictStr] = None """Default catalog for UC feature tables. Overrides ``DatabricksUCOfflineStoreConfig.default_catalog``.""" - schema: Optional[StrictStr] = None + uc_schema: Optional[StrictStr] = None """Default schema for UC feature tables. Overrides ``DatabricksUCOfflineStoreConfig.default_schema``.""" diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py index 87ee8b1a2b9..05b424073c4 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py @@ -89,15 +89,17 @@ def _feast_to_spark_type_simple(field): "STRING": StringType(), "UTF8": StringType(), "BINARY": BinaryType(), + "BYTES": BinaryType(), "TIMESTAMP": TimestampType(), "TIMESTAMP_TZ": TimestampType(), + "UNIXTIMESTAMP": TimestampType(), "DATE": DateType(), "DATE32": DateType(), } if type_name.startswith("LIST<") or type_name.startswith("ARRAY<"): return ArrayType(StringType()) - if "DOUBLE" in type_name or "FLOAT" in type_name and "64" in type_name: + if "DOUBLE" in type_name or ("FLOAT" in type_name and "64" in type_name): return DoubleType() if "LIST" in type_name or "ARRAY" in type_name: return ArrayType(StringType()) @@ -114,7 +116,7 @@ def _resolve_uc_path( fv: FeatureView, default_catalog: Optional[str], default_schema: Optional[str], -) -> Tuple[str, str, str]: +) -> Tuple[Optional[str], Optional[str], str]: """Resolve the (catalog, schema, table_name) for a feature view. Prioritises per‑FV tag overrides, then global defaults, then the @@ -364,7 +366,7 @@ def register_uc_feature_tables( fe_client = FeatureEngineeringClient() default_catalog = uc_config.catalog or config.default_catalog - default_schema = uc_config.schema or config.default_schema + default_schema = uc_config.uc_schema or config.default_schema for fv in feature_views: try: diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py new file mode 100644 index 00000000000..c394b1c1353 --- /dev/null +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py @@ -0,0 +1,448 @@ +from unittest.mock import MagicMock, patch + +from pyspark.sql.types import ( + ArrayType, + BinaryType, + BooleanType, + DoubleType, + FloatType, + IntegerType, + LongType, + StringType, + StructType, + TimestampType, +) + +from feast import FeatureView +from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( + DatabricksUCOfflineStoreConfig, + UCRegistrationConfig, +) +from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( + _build_full_table_name, + _build_spark_schema, + _build_uc_tags, + _feast_to_spark_type_simple, + _get_primary_keys, + _resolve_uc_path, + _should_register, + register_uc_feature_tables, +) + + +def make_mock_field(name: str, dtype_str: str = "", nullable: bool = True): + field = MagicMock() + field.name = name + # _feast_to_spark_type_simple does str(dtype).upper() + # When dtype_str is empty, dtype will be None -> StringType fallback + if dtype_str: + + class FakeDtype: + def __str__(self): + return dtype_str + + field.dtype = FakeDtype() + else: + field.dtype = None + return field + + +def test_feast_to_spark_type_simple_none(): + field = make_mock_field("col", "") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, StringType) + + +def test_feast_to_spark_type_simple_int32(): + field = make_mock_field("col", "Int32") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, IntegerType) + + +def test_feast_to_spark_type_simple_int64(): + field = make_mock_field("col", "Int64") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, LongType) + + +def test_feast_to_spark_type_simple_float32(): + field = make_mock_field("col", "Float32") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, FloatType) + + +def test_feast_to_spark_type_simple_float64(): + field = make_mock_field("col", "Float64") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, DoubleType) + + +def test_feast_to_spark_type_simple_string(): + field = make_mock_field("col", "String") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, StringType) + + +def test_feast_to_spark_type_simple_bool(): + field = make_mock_field("col", "Bool") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, BooleanType) + + +def test_feast_to_spark_type_simple_bytes(): + field = make_mock_field("col", "Bytes") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, BinaryType) + + +def test_feast_to_spark_type_simple_unix_timestamp(): + field = make_mock_field("col", "UnixTimestamp") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, TimestampType) + + +def test_feast_to_spark_type_simple_list(): + field = make_mock_field("col", "List") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, ArrayType) + + +def test_feast_to_spark_type_simple_array(): + field = make_mock_field("col", "Array") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, ArrayType) + + +def test_feast_to_spark_type_simple_unknown(): + field = make_mock_field("col", "UnknownType") + result = _feast_to_spark_type_simple(field) + assert isinstance(result, StringType) + + +def test_should_register_default(): + fv = MagicMock(spec=FeatureView) + fv.tags = {} + assert _should_register(fv) is True + + +def test_should_register_true(): + fv = MagicMock(spec=FeatureView) + fv.tags = {"uc.register_as_feature_table": "true"} + assert _should_register(fv) is True + + +def test_should_register_false(): + fv = MagicMock(spec=FeatureView) + fv.tags = {"uc.register_as_feature_table": "false"} + assert _should_register(fv) is False + + +def test_should_register_case_insensitive(): + fv = MagicMock(spec=FeatureView) + fv.tags = {"uc.register_as_feature_table": "FALSE"} + assert _should_register(fv) is False + + +def test_resolve_uc_path_uses_defaults(): + fv = MagicMock(spec=FeatureView) + fv.tags = {} + fv.name = "my_feature_view" + catalog, schema, table = _resolve_uc_path(fv, "default_cat", "default_sch") + assert catalog == "default_cat" + assert schema == "default_sch" + assert table == "my_feature_view" + + +def test_resolve_uc_path_uses_tags(): + fv = MagicMock(spec=FeatureView) + fv.tags = { + "uc.catalog": "tag_cat", + "uc.schema": "tag_sch", + "uc.table": "tag_table", + } + fv.name = "ignored" + catalog, schema, table = _resolve_uc_path(fv, "default_cat", "default_sch") + assert catalog == "tag_cat" + assert schema == "tag_sch" + assert table == "tag_table" + + +def test_resolve_uc_path_sanitizes_table_name(): + fv = MagicMock(spec=FeatureView) + fv.tags = {} + fv.name = "my-feature.view" + catalog, schema, table = _resolve_uc_path(fv, "cat", "sch") + assert table == "my_feature_view" + + +def test_resolve_uc_path_tag_overrides_none_default(): + fv = MagicMock(spec=FeatureView) + fv.tags = {"uc.catalog": "tag_cat"} + fv.name = "fv_name" + catalog, schema, table = _resolve_uc_path(fv, None, None) + assert catalog == "tag_cat" + assert schema is None + assert table == "fv_name" + + +def test_build_uc_tags(): + fv = MagicMock(spec=FeatureView) + fv.tags = { + "env": "prod", + "uc.register_as_feature_table": "false", + "owner": "team_a", + } + tags = _build_uc_tags(fv, "my_project") + assert tags["env"] == "prod" + assert tags["owner"] == "team_a" + assert "uc.register_as_feature_table" not in tags + assert tags["feast_managed"] == "feast" + assert tags["feast_project"] == "my_project" + + +def test_build_uc_tags_empty_tags(): + fv = MagicMock(spec=FeatureView) + fv.tags = {} + tags = _build_uc_tags(fv, "proj") + assert tags["feast_managed"] == "feast" + assert tags["feast_project"] == "proj" + + +def test_get_primary_keys(): + fv = MagicMock(spec=FeatureView) + col1 = make_mock_field("id", "Int32") + col2 = make_mock_field("date", "String") + fv.entity_columns = [col1, col2] + assert _get_primary_keys(fv) == ["id", "date"] + + +def test_get_primary_keys_empty(): + fv = MagicMock(spec=FeatureView) + fv.entity_columns = [] + assert _get_primary_keys(fv) == [] + + +def test_build_full_table_name_all_parts(): + assert _build_full_table_name("cat", "sch", "table") == "`cat`.`sch`.`table`" + + +def test_build_full_table_name_no_catalog(): + assert _build_full_table_name(None, "sch", "table") == "`sch`.`table`" + + +def test_build_full_table_name_no_schema(): + assert _build_full_table_name("cat", None, "table") == "`cat`.`table`" + + +def test_build_full_table_name_only_table(): + assert _build_full_table_name(None, None, "table") == "`table`" + + +def test_build_spark_schema(): + fv = MagicMock(spec=FeatureView) + id_field = make_mock_field("entity_id", "Int32") + feat_field = make_mock_field("feature_val", "Float64") + fv.entity_columns = [id_field] + fv.features = [feat_field] + fv.batch_source = MagicMock() + fv.batch_source.timestamp_field = "event_ts" + + schema = _build_spark_schema(fv) + + assert isinstance(schema, StructType) + fields = schema.fields + assert fields[0].name == "entity_id" + assert isinstance(fields[0].dataType, IntegerType) + assert fields[0].nullable is False + + assert fields[1].name == "feature_val" + assert isinstance(fields[1].dataType, DoubleType) + assert fields[1].nullable is True + + assert fields[2].name == "event_ts" + assert isinstance(fields[2].dataType, TimestampType) + assert fields[2].nullable is True + + +def test_build_spark_schema_no_timestamp(): + fv = MagicMock(spec=FeatureView) + id_field = make_mock_field("id", "Int64") + fv.entity_columns = [id_field] + fv.features = [] + fv.batch_source = MagicMock() + fv.batch_source.timestamp_field = None + + schema = _build_spark_schema(fv) + assert len(schema.fields) == 1 + assert schema.fields[0].name == "id" + + +def test_register_uc_feature_tables_skips_when_disabled(): + config = DatabricksUCOfflineStoreConfig(type="databricks_uc") + register_uc_feature_tables(config, [], "test_project") + # No error means success (skip silently) + + +def test_register_uc_feature_tables_skips_when_no_uc_config(): + config = DatabricksUCOfflineStoreConfig( + type="databricks_uc", + uc_registration=UCRegistrationConfig(enabled=False), + ) + register_uc_feature_tables(config, [], "test_project") + # No error means success + + +def _run_with_mock_fe(config, fvs, project, table_exists: bool = False): + """Helper to run register_uc_feature_tables with mocked external deps.""" + fe_module = MagicMock() + fe_module.FeatureEngineeringClient = MagicMock() + with patch.dict("sys.modules", {"databricks.feature_engineering": fe_module}): + with patch( + "feast.infra.offline_stores.contrib.spark_offline_store.uc_registration.get_databricks_session" + ) as mock_get_session: + fe_client = MagicMock() + fe_module.FeatureEngineeringClient.return_value = fe_client + mock_session = MagicMock() + mock_session.catalog.tableExists.return_value = table_exists + mock_get_session.return_value = mock_session + register_uc_feature_tables(config, fvs, project) + return fe_client, mock_session + + +def test_register_uc_feature_tables_creates_new_table(): + config = DatabricksUCOfflineStoreConfig( + type="databricks_uc", + default_catalog="cat", + default_schema="sch", + uc_registration=UCRegistrationConfig(enabled=True), + ) + + fv = MagicMock(spec=FeatureView) + fv.tags = {} + fv.name = "test_fv" + fv.entity_columns = [] + fv.features = [] + fv.batch_source = MagicMock() + fv.batch_source.timestamp_field = None + fv.description = "test description" + fv.owner = "user1" + + fe_client, mock_session = _run_with_mock_fe( + config, [fv], "proj", table_exists=False + ) + + fe_client.create_table.assert_called_once() + call_kwargs = fe_client.create_table.call_args[1] + assert call_kwargs["name"] == "`cat`.`sch`.`test_fv`" + assert call_kwargs["primary_keys"] == [] + assert call_kwargs["description"] == "test description" + + +@patch( + "feast.infra.offline_stores.contrib.spark_offline_store.uc_registration.get_databricks_session" +) +def test_register_uc_feature_tables_updates_existing( + mock_get_session, +): + mock_session = MagicMock() + mock_get_session.return_value = mock_session + mock_session.catalog.tableExists.return_value = True + + fe_module = MagicMock() + fe_client = MagicMock() + fe_module.FeatureEngineeringClient.return_value = fe_client + + config = DatabricksUCOfflineStoreConfig( + type="databricks_uc", + default_catalog="cat", + default_schema="sch", + uc_registration=UCRegistrationConfig(enabled=True), + ) + + fv = MagicMock(spec=FeatureView) + fv.tags = {} + fv.name = "test_fv" + fv.entity_columns = [] + fv.features = [] + fv.batch_source = MagicMock() + fv.batch_source.timestamp_field = None + fv.description = "test description" + fv.owner = "user1" + + with patch.dict("sys.modules", {"databricks.feature_engineering": fe_module}): + register_uc_feature_tables(config, [fv], "proj") + + fe_client.create_table.assert_not_called() + mock_session.sql.assert_any_call( + "COMMENT ON TABLE `cat`.`sch`.`test_fv` IS 'test description'" + ) + + +@patch( + "feast.infra.offline_stores.contrib.spark_offline_store.uc_registration.get_databricks_session" +) +def test_register_uc_feature_tables_skips_opt_out( + mock_get_session, +): + mock_session = MagicMock() + mock_get_session.return_value = mock_session + mock_session.catalog.tableExists.return_value = False + + config = DatabricksUCOfflineStoreConfig( + type="databricks_uc", + default_catalog="cat", + default_schema="sch", + uc_registration=UCRegistrationConfig(enabled=True), + ) + + fv = MagicMock(spec=FeatureView) + fv.tags = {"uc.register_as_feature_table": "false"} + fv.name = "opt_out_fv" + fv.entity_columns = [] + fv.features = [] + fv.batch_source = MagicMock() + fv.batch_source.timestamp_field = None + fv.description = "" + fv.owner = None + + fe_module = MagicMock() + fe_client = MagicMock() + fe_module.FeatureEngineeringClient.return_value = fe_client + with patch.dict("sys.modules", {"databricks.feature_engineering": fe_module}): + register_uc_feature_tables(config, [fv], "proj") + + # Per-view opt-out means create_table should NOT be called + fe_client.create_table.assert_not_called() + + +def test_register_uc_feature_tables_skips_when_missing_catalog(): + fe_module = MagicMock() + fe_client = MagicMock() + fe_module.FeatureEngineeringClient.return_value = fe_client + with patch.dict("sys.modules", {"databricks.feature_engineering": fe_module}): + with patch( + "feast.infra.offline_stores.contrib.spark_offline_store.uc_registration.get_databricks_session" + ) as mock_get_session: + mock_session = MagicMock() + mock_get_session.return_value = mock_session + + config = DatabricksUCOfflineStoreConfig( + type="databricks_uc", + default_catalog=None, + default_schema=None, + uc_registration=UCRegistrationConfig(enabled=True), + ) + + fv = MagicMock(spec=FeatureView) + fv.tags = {} + fv.name = "no_cat_fv" + fv.entity_columns = [] + fv.features = [] + fv.batch_source = MagicMock() + fv.batch_source.timestamp_field = None + fv.description = "" + fv.owner = None + + register_uc_feature_tables(config, [fv], "proj") + + fe_client.create_table.assert_not_called() From 16870f37d960bb5a85df0a4ec346b9f0c3f91d01 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Wed, 1 Jul 2026 17:31:39 +0530 Subject: [PATCH 06/16] feat: Refactor L2 to build on Iceberg REST Catalog instead of Databricks-specific SDKs Signed-off-by: Abhishek Shinde --- sdk/python/feast/feature_store.py | 8 +- .../spark_offline_store/databricks_uc.py | 107 +++-- .../contrib/spark_offline_store/spark.py | 87 ++-- .../spark_offline_store/uc_registration.py | 389 ------------------ .../contrib/spark_offline_store/utils.py | 101 +++++ .../infra/offline_stores/iceberg/__init__.py | 0 .../offline_stores/iceberg/catalog_config.py | 61 +++ .../offline_stores/iceberg/catalog_manager.py | 371 +++++++++++++++++ .../offline_stores/iceberg/iceberg_source.py | 352 ++++++++++++++++ .../offline_stores/iceberg/registration.py | 188 +++++++++ .../spark_offline_store/test_databricks_uc.py | 82 +++- .../test_iceberg_source.py | 210 ++++++++++ .../test_uc_registration.py | 366 ++++++---------- 13 files changed, 1600 insertions(+), 722 deletions(-) delete mode 100644 sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py create mode 100644 sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py create mode 100644 sdk/python/feast/infra/offline_stores/iceberg/__init__.py create mode 100644 sdk/python/feast/infra/offline_stores/iceberg/catalog_config.py create mode 100644 sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py create mode 100644 sdk/python/feast/infra/offline_stores/iceberg/iceberg_source.py create mode 100644 sdk/python/feast/infra/offline_stores/iceberg/registration.py create mode 100644 sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_iceberg_source.py diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index b143b41a595..701027d7537 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -1688,7 +1688,7 @@ def _register_uc_feature_tables_from_diffs( from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( DatabricksUCOfflineStoreConfig, ) - from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( + from feast.infra.offline_stores.iceberg.registration import ( register_uc_feature_tables, ) @@ -1707,7 +1707,7 @@ def _register_uc_feature_tables_from_diffs( if not fvs: return - register_uc_feature_tables(self.config.offline_store, fvs, self.project) + register_uc_feature_tables(self.config.offline_store.catalog, fvs, self.project) def _register_uc_feature_tables_legacy(self, objects: List[Any]) -> None: """Register feature views as UC feature tables (legacy apply path). @@ -1718,7 +1718,7 @@ def _register_uc_feature_tables_legacy(self, objects: List[Any]) -> None: from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( DatabricksUCOfflineStoreConfig, ) - from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( + from feast.infra.offline_stores.iceberg.registration import ( register_uc_feature_tables, ) @@ -1733,7 +1733,7 @@ def _register_uc_feature_tables_legacy(self, objects: List[Any]) -> None: if not fvs: return - register_uc_feature_tables(self.config.offline_store, fvs, self.project) + register_uc_feature_tables(self.config.offline_store.catalog, fvs, self.project) def teardown(self): """Tears down all local and cloud resources for the feature store.""" diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py index 3fd48b67d96..7eacd5a4481 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py @@ -15,6 +15,7 @@ SparkOfflineStore, SparkOfflineStoreConfig, ) +from feast.infra.offline_stores.iceberg.catalog_config import IcebergCatalogConfig from feast.infra.offline_stores.offline_store import RetrievalJob from feast.infra.registry.base_registry import BaseRegistry from feast.repo_config import FeastConfigBaseModel, RepoConfig @@ -26,33 +27,38 @@ class UCRegistrationConfig(FeastConfigBaseModel): """Configuration for Unity Catalog feature table registration during ``feast apply``.""" enabled: StrictBool = True - """Whether to register feature views as UC feature tables on ``feast apply``.""" + """Whether to register feature views as catalog feature tables on ``feast apply``.""" catalog: Optional[StrictStr] = None - """Default catalog for UC feature tables. Overrides ``DatabricksUCOfflineStoreConfig.default_catalog``.""" + """Default catalog for feature tables. Overrides ``IcebergCatalogConfig.warehouse``.""" uc_schema: Optional[StrictStr] = None - """Default schema for UC feature tables. Overrides ``DatabricksUCOfflineStoreConfig.default_schema``.""" + """Default schema for feature tables. Overrides ``IcebergCatalogConfig.namespace``.""" class DatabricksUCOfflineStoreConfig(SparkOfflineStoreConfig): - type: StrictStr = "databricks_uc" - """Offline store type selector""" + """Offline store configuration for Databricks Unity Catalog. - workspace_host: Optional[StrictStr] = None - """Databricks workspace host (e.g. adb-xxxx.azuredatabricks.net)""" + Uses an Iceberg REST Catalog connection to resolve table metadata + and a Spark session (via Databricks Connect or standard Spark) to + read and write data. + """ - token: Optional[StrictStr] = None - """Databricks Personal Access Token (PAT)""" + type: StrictStr = "databricks_uc" + """Offline store type selector""" - cluster_id: Optional[StrictStr] = None - """Databricks Cluster ID to connect to for Databricks Connect""" + catalog: IcebergCatalogConfig + """Iceberg REST Catalog connection configuration. - default_catalog: Optional[StrictStr] = None - """Default catalog name to use in Unity Catalog""" + For Databricks UC:: - default_schema: Optional[StrictStr] = None - """Default schema name to use in Unity Catalog""" + catalog: + type: rest + endpoint: https://.databricks.net/api/2.1/unity-catalog/iceberg + warehouse: prod_ml + namespace: features + token_env_var: DATABRICKS_TOKEN + """ uc_registration: Optional[UCRegistrationConfig] = None """Configuration for UC feature table registration during ``feast apply``.""" @@ -61,27 +67,30 @@ class DatabricksUCOfflineStoreConfig(SparkOfflineStoreConfig): def get_databricks_session( store_config: DatabricksUCOfflineStoreConfig, ) -> SparkSession: - # Check if there is already an active session + """Create (or reuse) a Databricks Spark session. + + Uses ``IcebergCatalogConfig`` to derive the Databricks workspace + host and token, then initializes a Databricks Connect session. + + Falls back to a regular Spark session when ``databricks-connect`` + is not installed. + """ spark_session = SparkSession.getActiveSession() if not spark_session: - workspace_host = store_config.workspace_host - token = store_config.token - cluster_id = store_config.cluster_id + catalog_cfg = store_config.catalog + endpoint = catalog_cfg.endpoint - # Clean host URL if it starts with https:// - if workspace_host: - if workspace_host.startswith("https://"): - workspace_host = workspace_host[8:] - elif workspace_host.startswith("http://"): - workspace_host = workspace_host[7:] + # Extract workspace host from the Iceberg REST endpoint URL. + # UC endpoints look like: + # https://.databricks.net/api/2.1/unity-catalog/iceberg + workspace_host = _extract_workspace_host(endpoint) + token = _resolve_token(catalog_cfg) - if workspace_host and cluster_id: - # Databricks Connect V2 initialization (Spark Connect URI format) + if workspace_host: conn_str = f"sc://{workspace_host}:443/" params = [] if token: params.append(f"token={token}") - params.append(f"x-databricks-cluster-id={cluster_id}") if params: conn_str = f"{conn_str};{';'.join(params)}" @@ -90,7 +99,6 @@ def get_databricks_session( builder = DatabricksSession.builder.remote(conn_str) except ImportError: - # Fallback to standard PySpark remote connect if databricks-connect not installed builder = SparkSession.builder.remote(conn_str) else: try: @@ -110,18 +118,48 @@ def get_databricks_session( assert spark_session is not None - # Apply configuration defaults spark_session.conf.set("spark.sql.parser.quotedRegexColumnNames", "true") - if store_config.default_catalog: - spark_session.sql(f"USE CATALOG `{store_config.default_catalog}`") - if store_config.default_schema: - spark_session.sql(f"USE SCHEMA `{store_config.default_schema}`") + if store_config.catalog.warehouse: + spark_session.sql(f"USE CATALOG `{store_config.catalog.warehouse}`") + if store_config.catalog.namespace: + spark_session.sql(f"USE SCHEMA `{store_config.catalog.namespace}`") return spark_session +def _extract_workspace_host(endpoint: str) -> Optional[str]: + """Extract the Databricks workspace host from an Iceberg REST endpoint URL.""" + if not endpoint: + return None + cleaned = endpoint + if cleaned.startswith("https://"): + cleaned = cleaned[8:] + elif cleaned.startswith("http://"): + cleaned = cleaned[7:] + # Stop at the first slash (path portion) + idx = cleaned.find("/") + if idx != -1: + cleaned = cleaned[:idx] + return cleaned if cleaned else None + + +def _resolve_token(catalog_cfg: IcebergCatalogConfig) -> str: + """Resolve the bearer token from the configured env var.""" + import os + + return os.environ.get(catalog_cfg.token_env_var, "") + + class DatabricksUCOfflineStore(SparkOfflineStore): + """Databricks Unity Catalog offline store. + + A thin wrapper around ``SparkOfflineStore`` that ensures a + Databricks Connect session is active before delegating to the + parent class. Table metadata is resolved via the Iceberg REST + Catalog (``IcebergCatalogConfig`` / ``pyiceberg``). + """ + @staticmethod def pull_latest_from_table_or_query( config: RepoConfig, @@ -134,7 +172,6 @@ def pull_latest_from_table_or_query( end_date: datetime, ) -> RetrievalJob: assert isinstance(config.offline_store, DatabricksUCOfflineStoreConfig) - # Initialize/Retrieve the Databricks Spark Session so it's registered as active get_databricks_session(config.offline_store) return SparkOfflineStore.pull_latest_from_table_or_query( diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py index 8c7e657b257..f7030417cbb 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py @@ -45,6 +45,10 @@ SavedDatasetSparkStorage, SparkSource, ) +from feast.infra.offline_stores.iceberg.iceberg_source import ( + IcebergRestCatalogSource, + UnityCatalogSource, +) from feast.infra.offline_stores.offline_store import ( OfflineStore, RetrievalJob, @@ -104,11 +108,11 @@ def pull_latest_from_table_or_query( start_date: datetime, end_date: datetime, ) -> RetrievalJob: - spark_session = get_spark_session_or_start_new_with_repoconfig( - config.offline_store + assert isinstance(data_source, (SparkSource, IcebergRestCatalogSource)), ( + f"SparkOfflineStore requires SparkSource or IcebergRestCatalogSource, got {type(data_source)}" ) - assert isinstance(config.offline_store, SparkOfflineStoreConfig) - assert isinstance(data_source, SparkSource) + + spark_session = _resolve_spark_session_for_source(data_source, config) warnings.warn( "The spark offline store is an experimental feature in alpha development. " @@ -174,10 +178,13 @@ def get_historical_features( full_feature_names: bool = False, **kwargs, ) -> RetrievalJob: - assert isinstance(config.offline_store, SparkOfflineStoreConfig) date_partition_column_formats = [] for fv in feature_views: - assert isinstance(fv.batch_source, SparkSource) + assert isinstance( + fv.batch_source, (SparkSource, IcebergRestCatalogSource) + ), ( + f"SparkOfflineStore requires SparkSource or IcebergRestCatalogSource batch sources, got {type(fv.batch_source)}" + ) date_partition_column_formats.append( fv.batch_source.date_partition_column_format ) @@ -188,8 +195,8 @@ def get_historical_features( RuntimeWarning, ) - spark_session = get_spark_session_or_start_new_with_repoconfig( - store_config=config.offline_store + spark_session = _resolve_spark_session_for_source( + feature_views[0].batch_source, config ) tmp_entity_df_table_name = offline_utils.get_temp_entity_table_name() @@ -329,8 +336,11 @@ def offline_write_batch( table: pyarrow.Table, progress: Optional[Callable[[int], Any]], ): - assert isinstance(config.offline_store, SparkOfflineStoreConfig) - assert isinstance(feature_view.batch_source, SparkSource) + assert isinstance( + feature_view.batch_source, (SparkSource, IcebergRestCatalogSource) + ), ( + f"SparkOfflineStore requires SparkSource or IcebergRestCatalogSource batch sources, got {type(feature_view.batch_source)}" + ) pa_schema, column_names = offline_utils.get_pyarrow_schema_from_batch_source( config, feature_view.batch_source @@ -341,8 +351,8 @@ def offline_write_batch( f"The schema is expected to be {pa_schema} with the columns (in this exact order) to be {column_names}." ) - spark_session = get_spark_session_or_start_new_with_repoconfig( - store_config=config.offline_store + spark_session = _resolve_spark_session_for_source( + feature_view.batch_source, config ) if feature_view.batch_source.path: @@ -390,17 +400,16 @@ def pull_all_from_table_or_query( created_timestamp_column have all already been mapped to column names of the source table and those column names are the values passed into this function. """ - assert isinstance(config.offline_store, SparkOfflineStoreConfig) - assert isinstance(data_source, SparkSource) + assert isinstance(data_source, (SparkSource, IcebergRestCatalogSource)), ( + f"SparkOfflineStore requires SparkSource or IcebergRestCatalogSource, got {type(data_source)}" + ) warnings.warn( "The spark offline store is an experimental feature in alpha development. " "This API is unstable and it could and most probably will be changed in the future.", RuntimeWarning, ) - spark_session = get_spark_session_or_start_new_with_repoconfig( - store_config=config.offline_store - ) + spark_session = _resolve_spark_session_for_source(data_source, config) timestamp_fields = [timestamp_field] if created_timestamp_column: @@ -447,12 +456,11 @@ def compute_monitoring_metrics( histogram_bins: int = 20, top_n: int = 10, ) -> List[Dict[str, Any]]: - assert isinstance(config.offline_store, SparkOfflineStoreConfig) - assert isinstance(data_source, SparkSource) - - spark_session = get_spark_session_or_start_new_with_repoconfig( - store_config=config.offline_store + assert isinstance(data_source, (SparkSource, IcebergRestCatalogSource)), ( + f"SparkOfflineStore requires SparkSource or IcebergRestCatalogSource, got {type(data_source)}" ) + + spark_session = _resolve_spark_session_for_source(data_source, config) from_expression = data_source.get_table_query_string() ts_filter = get_timestamp_filter_sql( start_date, @@ -497,12 +505,11 @@ def get_monitoring_max_timestamp( data_source: DataSource, timestamp_field: str, ) -> Optional[datetime]: - assert isinstance(config.offline_store, SparkOfflineStoreConfig) - assert isinstance(data_source, SparkSource) - - spark_session = get_spark_session_or_start_new_with_repoconfig( - store_config=config.offline_store + assert isinstance(data_source, (SparkSource, IcebergRestCatalogSource)), ( + f"SparkOfflineStore requires SparkSource or IcebergRestCatalogSource, got {type(data_source)}" ) + + spark_session = _resolve_spark_session_for_source(data_source, config) from_expression = data_source.get_table_query_string() q_ts = f"`{timestamp_field}`" sql = f"SELECT MAX({q_ts}) AS max_ts FROM {from_expression} AS _src" @@ -1176,6 +1183,28 @@ def get_spark_session_or_start_new_with_repoconfig( return spark_session +def _resolve_spark_session_for_source( + data_source: DataSource, + config: RepoConfig, +) -> SparkSession: + """Return an appropriate Spark session based on the data source type. + + - For ``UnityCatalogSource``: create a Databricks Connect session. + - For plain ``IcebergRestCatalogSource``: fall back to a standard Spark session. + - For ``SparkSource``: use the existing ``get_spark_session_or_start_new_with_repoconfig``. + """ + if isinstance(data_source, UnityCatalogSource): + from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( + get_databricks_session, + ) + + return get_databricks_session(config.offline_store) + elif isinstance(data_source, IcebergRestCatalogSource): + return get_spark_session_or_start_new_with_repoconfig(config.offline_store) + else: + return get_spark_session_or_start_new_with_repoconfig(config.offline_store) + + def _gather_all_entities( fv_query_contexts: List[offline_utils.FeatureViewQueryContext], ) -> List[str]: @@ -1208,7 +1237,9 @@ def _create_temp_entity_union_view( for fv, ctx, date_format in zip( feature_views, fv_query_contexts, date_partition_column_formats ): - assert isinstance(fv.batch_source, SparkSource) + assert isinstance(fv.batch_source, (SparkSource, IcebergRestCatalogSource)), ( + f"SparkOfflineStore requires SparkSource or IcebergRestCatalogSource batch sources, got {type(fv.batch_source)}" + ) from_expression = fv.batch_source.get_table_query_string() timestamp_field = fv.batch_source.timestamp_field or "event_timestamp" date_partition_column = fv.batch_source.date_partition_column diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py deleted file mode 100644 index 05b424073c4..00000000000 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/uc_registration.py +++ /dev/null @@ -1,389 +0,0 @@ -"""Unity Catalog feature table registration for ``feast apply``. - -When the offline store is configured as ``databricks_uc``, this module -registers (or updates) each FeatureView as a Unity Catalog feature table -via the Databricks ``FeatureEngineeringClient``. FeatureView entities -become UC primary keys, and tags/description/owner are synced to UC metadata. - -Per‑FeatureView opt‑out and overrides are controlled via FeatureView ``tags``: - -* ``uc.register_as_feature_table`` — ``"false"`` skips a specific view. -* ``uc.catalog`` / ``uc.schema`` / ``uc.table`` — override the UC path. - -Global defaults come from ``UCRegistrationConfig`` (inside -``DatabricksUCOfflineStoreConfig``). -""" - -import logging -from typing import Dict, List, Optional, Tuple - -import click -from pyspark.sql import SparkSession -from pyspark.sql.types import ( - StructField, - StructType, -) - -from feast import FeatureView -from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( - DatabricksUCOfflineStoreConfig, - get_databricks_session, -) - -logger = logging.getLogger(__name__) - -# FeatureView tag keys for per‑FV UC configuration -_REGISTER_AS_FEATURE_TABLE_KEY = "uc.register_as_feature_table" -_CATALOG_KEY = "uc.catalog" -_SCHEMA_KEY = "uc.schema" -_TABLE_KEY = "uc.table" - -# Internal Feast tags stored on the UC table -_MANAGED_BY_TAG = "feast_managed" - - -def _feast_to_spark_type_simple(field): - """Convert a Feast :class:`Field` to a pyspark :class:`DataType`. - - This is a best‑effort mapping; the underlying Spark‑read path performs - full schema inference at query time. - """ - from pyspark.sql.types import ( - ArrayType, - BinaryType, - BooleanType, - ByteType, - DateType, - DoubleType, - FloatType, - IntegerType, - LongType, - ShortType, - StringType, - TimestampType, - ) - - dtype = getattr(field, "dtype", None) - if dtype is None: - return StringType() # fallback - - type_name = str(dtype).upper() - - type_map = { - "BOOL": BooleanType(), - "BOOLEAN": BooleanType(), - "INT8": ByteType(), - "BYTE": ByteType(), - "INT16": ShortType(), - "SHORT": ShortType(), - "INT32": IntegerType(), - "INT": IntegerType(), - "INTEGER": IntegerType(), - "INT64": LongType(), - "LONG": LongType(), - "BIGINT": LongType(), - "FLOAT32": FloatType(), - "FLOAT": FloatType(), - "FLOAT64": DoubleType(), - "DOUBLE": DoubleType(), - "STRING": StringType(), - "UTF8": StringType(), - "BINARY": BinaryType(), - "BYTES": BinaryType(), - "TIMESTAMP": TimestampType(), - "TIMESTAMP_TZ": TimestampType(), - "UNIXTIMESTAMP": TimestampType(), - "DATE": DateType(), - "DATE32": DateType(), - } - - if type_name.startswith("LIST<") or type_name.startswith("ARRAY<"): - return ArrayType(StringType()) - if "DOUBLE" in type_name or ("FLOAT" in type_name and "64" in type_name): - return DoubleType() - if "LIST" in type_name or "ARRAY" in type_name: - return ArrayType(StringType()) - - return type_map.get(type_name, StringType()) - - -def _should_register(fv: FeatureView) -> bool: - """Return ``True`` unless the feature view opts out via tags.""" - return fv.tags.get(_REGISTER_AS_FEATURE_TABLE_KEY, "true").lower() != "false" - - -def _resolve_uc_path( - fv: FeatureView, - default_catalog: Optional[str], - default_schema: Optional[str], -) -> Tuple[Optional[str], Optional[str], str]: - """Resolve the (catalog, schema, table_name) for a feature view. - - Prioritises per‑FV tag overrides, then global defaults, then the - feature view name as the table name. - """ - catalog = fv.tags.get(_CATALOG_KEY) or default_catalog - schema = fv.tags.get(_SCHEMA_KEY) or default_schema - table = fv.tags.get(_TABLE_KEY) or fv.name - - # Sanitise: replace characters that are invalid in UC names - table = table.replace("-", "_").replace(".", "_").replace(" ", "_") - return catalog, schema, table - - -def _build_spark_schema(fv: FeatureView) -> StructType: - """Build a pyspark ``StructType`` from the feature view's columns.""" - fields: List[StructField] = [] - - seen: set = set() - for col in fv.entity_columns: - if col.name not in seen: - fields.append( - StructField(col.name, _feast_to_spark_type_simple(col), nullable=False) - ) - seen.add(col.name) - - for col in fv.features: - if col.name not in seen: - fields.append( - StructField(col.name, _feast_to_spark_type_simple(col), nullable=True) - ) - seen.add(col.name) - - # Add timestamp column if the batch source declares one - timestamp_field = getattr(fv.batch_source, "timestamp_field", None) - if timestamp_field and timestamp_field not in seen: - from pyspark.sql.types import TimestampType - - fields.append(StructField(timestamp_field, TimestampType(), nullable=True)) - - return StructType(fields) - - -def _get_primary_keys(fv: FeatureView) -> List[str]: - """Extract primary key column names from entity columns.""" - return [col.name for col in fv.entity_columns] - - -def _build_uc_tags(fv: FeatureView, project: str) -> Dict[str, str]: - """Build UC table tags, excluding internal ``uc.*`` keys.""" - tags: Dict[str, str] = {} - for key, value in fv.tags.items(): - if not key.startswith("uc."): - tags[key] = value - tags[_MANAGED_BY_TAG] = "feast" - tags["feast_project"] = project - return tags - - -def _escape_sql_string(value: str) -> str: - """Escape single quotes for SQL string literals.""" - return value.replace("\\", "\\\\").replace("'", "\\'") - - -def _build_full_table_name(catalog, schema, table): - """Build a three-level table reference: ``catalog.schema.table``.""" - parts = [] - if catalog: - parts.append(f"`{catalog}`") - if schema: - parts.append(f"`{schema}`") - parts.append(f"`{table}`") - return ".".join(parts) - - -def _table_exists(spark_session: SparkSession, full_name: str) -> bool: - """Check whether a UC table exists.""" - try: - return spark_session.catalog.tableExists(full_name) - except Exception: - return False - - -def _create_uc_feature_table( - fe_client, - spark_session: SparkSession, - fv: FeatureView, - full_name: str, - primary_keys: List[str], - project: str, -) -> None: - """Create a new UC feature table via FeatureEngineeringClient.""" - spark_schema = _build_spark_schema(fv) - timestamp_key = getattr(fv.batch_source, "timestamp_field", None) - description = fv.description or "" - tags = _build_uc_tags(fv, project) - owner = fv.owner - - fe_client.create_table( - name=full_name, - primary_keys=primary_keys, - schema=spark_schema, - timestamp_key=timestamp_key, - description=description, - tags=tags, - ) - - if owner: - try: - spark_session.sql( - f"ALTER TABLE {full_name} SET OWNER TO `{_escape_sql_string(owner)}`" - ) - except Exception: - logger.debug( - "Could not set owner for UC table %s; continuing", - full_name, - ) - - -def _update_uc_feature_table_metadata( - spark_session: SparkSession, - fv: FeatureView, - full_name: str, - project: str, -) -> None: - """Update metadata on an existing UC feature table.""" - description = fv.description or "" - owner = fv.owner - tags = _build_uc_tags(fv, project) - - # Update description (COMMENT) - try: - spark_session.sql( - f"COMMENT ON TABLE {full_name} IS '{_escape_sql_string(description)}'" - ) - except Exception: - logger.debug( - "Could not update comment on UC table %s; continuing", - full_name, - ) - - # Update tags - if tags: - tag_pairs = ", ".join( - f"'{_escape_sql_string(k)}' = '{_escape_sql_string(v)}'" - for k, v in tags.items() - ) - try: - spark_session.sql(f"ALTER TABLE {full_name} SET TAGS ({tag_pairs})") - except Exception: - logger.debug( - "Could not update tags on UC table %s; continuing", - full_name, - ) - - # Update owner - if owner: - try: - spark_session.sql( - f"ALTER TABLE {full_name} SET OWNER TO `{_escape_sql_string(owner)}`" - ) - except Exception: - logger.debug( - "Could not update owner on UC table %s; continuing", - full_name, - ) - - -def _register_single_feature_view( - spark_session: SparkSession, - fe_client, - fv: FeatureView, - project: str, - default_catalog: Optional[str], - default_schema: Optional[str], -) -> None: - """Register or update a single FeatureView as a UC feature table.""" - if not _should_register(fv): - click.echo(f" ⊘ Skipping UC registration for {fv.name} (opt‑out)") - return - - catalog, schema, table = _resolve_uc_path(fv, default_catalog, default_schema) - if not catalog or not schema: - click.echo( - f" ⚠ Cannot register {fv.name}: missing catalog or schema. " - f"Set default_catalog/default_schema in feature_store.yaml or " - f"use tags uc.catalog/uc.schema." - ) - return - - full_name = _build_full_table_name(catalog, schema, table) - primary_keys = _get_primary_keys(fv) - - if _table_exists(spark_session, full_name): - _update_uc_feature_table_metadata(spark_session, fv, full_name, project) - click.echo(f" ✓ Updated UC feature table: {full_name}") - else: - _create_uc_feature_table( - fe_client, - spark_session, - fv, - full_name, - primary_keys, - project, - ) - click.echo(f" ✓ Created UC feature table: {full_name}") - - -def register_uc_feature_tables( - config: DatabricksUCOfflineStoreConfig, - feature_views: List[FeatureView], - project: str, -) -> None: - """Register or update FeatureViews as Unity Catalog feature tables. - - Skips silently when: - - ``uc_registration`` is not configured or ``enabled`` is ``False``. - - ``databricks-feature-engineering`` is not installed. - - The Databricks Spark session cannot be created. - - Per‑feature‑view errors are logged and do not halt the apply. - """ - uc_config = config.uc_registration - if uc_config is None or not uc_config.enabled: - return - - try: - from databricks.feature_engineering import ( - FeatureEngineeringClient, # noqa: F401 - ) - except ImportError: - logger.info( - "databricks-feature-engineering is not installed; " - "skipping UC feature table registration. " - "Install with: pip install databricks-feature-engineering" - ) - return - - try: - spark_session = get_databricks_session(config) - except Exception as e: - logger.warning( - "Could not create Databricks Spark session for UC registration: %s", e - ) - return - - fe_client = FeatureEngineeringClient() - - default_catalog = uc_config.catalog or config.default_catalog - default_schema = uc_config.uc_schema or config.default_schema - - for fv in feature_views: - try: - _register_single_feature_view( - spark_session, - fe_client, - fv, - project, - default_catalog, - default_schema, - ) - except Exception: - logger.exception( - "Failed to register UC feature table for FeatureView '%s'", - fv.name, - ) - click.echo( - f" ✗ Failed to register UC feature table: {fv.name} " - f"(check logs for details)" - ) diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py new file mode 100644 index 00000000000..b797bd2a6c9 --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py @@ -0,0 +1,101 @@ +import logging +from typing import Dict, Optional + +from pyspark import SparkConf +from pyspark.sql import SparkSession + +logger = logging.getLogger(__name__) + + +def _normalise_host(host: str) -> str: + if host.startswith("https://"): + return host[8:] + if host.startswith("http://"): + return host[7:] + return host + + +def _expected_connect_uri(host: str) -> str: + return f"sc://{_normalise_host(host)}:443/" + + +def get_databricks_connect_session( + host: Optional[str] = None, + token: Optional[str] = None, + cluster_id: Optional[str] = None, + catalog: Optional[str] = None, + schema: Optional[str] = None, + spark_conf: Optional[Dict[str, str]] = None, +) -> SparkSession: + """Create (or reuse) a Databricks Connect Spark session. + + Args: + host: Databricks workspace host (e.g. ``adb-xxxx.azuredatabricks.net``). + token: Databricks PAT. + cluster_id: Databricks cluster ID for the remote session. + catalog: Default Unity Catalog to use after connecting. + schema: Default schema to use after connecting. + spark_conf: Additional Spark configuration. + + Returns: + A Spark session connected to the Databricks workspace. + """ + spark_session = SparkSession.getActiveSession() + if spark_session and host: + try: + current_uri = spark_session.conf.get("spark.connect.client.uri", "") + expected_uri = _expected_connect_uri(host) + if expected_uri not in current_uri: + logger.info( + "Active Spark session URI %r does not match requested host %s " + "(expected %r). Stopping and recreating.", + current_uri, + host, + expected_uri, + ) + spark_session.stop() + spark_session = None + except Exception: + spark_session = None + + if not spark_session: + if host and cluster_id: + conn_str = _expected_connect_uri(host) + params = [] + if token: + params.append(f"token={token}") + params.append(f"x-databricks-cluster-id={cluster_id}") + if params: + conn_str = f"{conn_str};{';'.join(params)}" + + try: + from databricks.connect import DatabricksSession + + builder = DatabricksSession.builder.remote(conn_str) + except ImportError: + builder = SparkSession.builder.remote(conn_str) + else: + try: + from databricks.connect import DatabricksSession + + builder = DatabricksSession.builder + except ImportError: + builder = SparkSession.builder + + if spark_conf: + builder = builder.config( + conf=SparkConf().setAll([(k, v) for k, v in spark_conf.items()]) + ) + + spark_session = builder.getOrCreate() + + assert spark_session is not None + + spark_session.conf.set("spark.sql.parser.quotedRegexColumnNames", "true") + + if catalog: + spark_session.sql(f"USE CATALOG `{catalog}`") + if schema: + spark_session.sql(f"USE SCHEMA `{schema}`") + + return spark_session diff --git a/sdk/python/feast/infra/offline_stores/iceberg/__init__.py b/sdk/python/feast/infra/offline_stores/iceberg/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/sdk/python/feast/infra/offline_stores/iceberg/catalog_config.py b/sdk/python/feast/infra/offline_stores/iceberg/catalog_config.py new file mode 100644 index 00000000000..3b5225e9a76 --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/iceberg/catalog_config.py @@ -0,0 +1,61 @@ +import logging +from typing import Optional + +from pydantic import StrictBool, StrictStr + +from feast.repo_config import FeastConfigBaseModel + +logger = logging.getLogger(__name__) + + +class IcebergCatalogConfig(FeastConfigBaseModel): + """Configuration for connecting to an Iceberg REST Catalog. + + This config is generic and works with any Iceberg REST catalog + implementation: Databricks Unity Catalog, Apache Polaris, Nessie, + Tabular, etc. Authentication is via an environment variable + (``token_env_var``) and the Iceberg REST ``v1/oauth/tokens`` + endpoint where supported. + + The three-level namespace (``warehouse.namespace.table``) maps to + the Iceberg REST Catalog's ``prefix`` (warehouse/prefix), + ``namespace`` (schema), and ``table`` concepts. + """ + + type: StrictStr = "rest" + """Catalog backend type. One of ``"rest"`` (generic REST catalog), + ``"databricks"``, ``"polaris"``, ``"nessie"``. + See ``pyiceberg.catalog.load_catalog()`` for the full list.""" + + endpoint: StrictStr + """Iceberg REST Catalog endpoint URL. + + For Databricks Unity Catalog this is typically: + ``https://.databricks.net/api/2.1/unity-catalog/iceberg`` + + For Apache Polaris this is the Polaris REST endpoint.""" + + warehouse: StrictStr + """Warehouse (sometimes called ``prefix`` or ``catalog`` in the + Iceberg REST protocol). For Unity Catalog this is the catalog name, + e.g. ``prod_ml``.""" + + namespace: StrictStr + """Default namespace (schema) within the warehouse, e.g. ``features``.""" + + token_env_var: StrictStr = "ICEBERG_REST_TOKEN" + """Environment variable name that holds the bearer token / PAT for + the Iceberg REST Catalog.""" + + credential_vending: StrictBool = True + """Whether to use the Iceberg REST ``v1/oauth/tokens`` credential + vending flow when configuring the underlying file IO for engines.""" + + warehouse_location: Optional[StrictStr] = None + """Optional storage location for warehouse data (``s3://...``, + ``abfss://...``, ``gs://...``). If not set, the catalog may + derive it from the endpoint response.""" + + rest_signing: StrictBool = False + """Enable REST sign‑and‑push for write operations when the REST + catalog supports the ``v1/configure`` signing endpoint.""" diff --git a/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py b/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py new file mode 100644 index 00000000000..1d4129595bb --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py @@ -0,0 +1,371 @@ +"""Thin wrapper around ``pyiceberg`` for feature‑store operations. + +All Iceberg REST Catalog interactions needed by Feast (table creation, +schema inspection, property management, and data writes) go through +this module so that the rest of the codebase never imports ``pyiceberg`` +directly. +""" + +import logging +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence + +from pyspark.sql import DataFrame as SparkDataFrame + +from feast.infra.offline_stores.iceberg.catalog_config import IcebergCatalogConfig + +if TYPE_CHECKING: + from pyiceberg.catalog import Catalog + from pyiceberg.schema import Schema + +logger = logging.getLogger(__name__) + +# Feast metadata stored as Iceberg table properties +FEAST_PRIMARY_KEYS_PROP = "feast.primary_keys" +FEAST_TIMESTAMP_KEY_PROP = "feast.timestamp_key" +FEAST_DESCRIPTION_PROP = "feast.description" +FEAST_OWNER_PROP = "feast.owner" +FEAST_PROJECT_PROP = "feast.project" +FEAST_MANAGED_PROP = "feast.managed" +_FEAST_TAG_PREFIX = "feast.tag." + + +def _build_catalog_props(config: IcebergCatalogConfig, token: str) -> Dict[str, str]: + """Build the property dict passed to ``pyiceberg.catalog.load_catalog()``.""" + import os + + props: Dict[str, str] = { + "uri": config.endpoint, + "warehouse": config.warehouse, + } + token_val = token or os.environ.get(config.token_env_var, "") + if token_val: + props["token"] = token_val + if config.credential_vending: + props["credential-vending"] = "true" + if config.rest_signing: + props["rest-signing"] = "true" + if config.warehouse_location: + props["warehouse-location"] = config.warehouse_location + return props + + +def load_catalog( + config: IcebergCatalogConfig, + token: str = "", +) -> "Catalog": + """Create an Iceberg ``Catalog`` from Feast configuration.""" + from pyiceberg.catalog import load_catalog as _load_catalog + + props = _build_catalog_props(config, token) + return _load_catalog(name=config.warehouse, type=config.type, **props) + + +def namespace_exists(catalog: "Catalog", namespace: str) -> bool: + """Return ``True`` if the namespace (schema) exists in the catalog.""" + try: + catalog.list_tables(namespace) + return True + except Exception: + return False + + +def table_exists(catalog: "Catalog", namespace: str, table_name: str) -> bool: + """Return ``True`` if the table exists in the catalog.""" + try: + catalog.load_table((namespace, table_name)) + return True + except Exception: + return False + + +def _feast_to_iceberg_type(feast_dtype) -> str: + """Convert a Feast ``dtype`` attribute to an Iceberg type string.""" + if feast_dtype is None: + return "string" + + type_name = str(feast_dtype).upper() + + type_map = { + "BOOL": "boolean", + "BOOLEAN": "boolean", + "INT8": "int", + "BYTE": "int", + "INT16": "int", + "SHORT": "int", + "INT32": "int", + "INT": "int", + "INTEGER": "int", + "INT64": "long", + "LONG": "long", + "BIGINT": "long", + "FLOAT32": "float", + "FLOAT": "float", + "FLOAT64": "double", + "DOUBLE": "double", + "STRING": "string", + "UTF8": "string", + "BINARY": "binary", + "BYTES": "binary", + "TIMESTAMP": "timestamp", + "TIMESTAMP_TZ": "timestamptz", + "UNIXTIMESTAMP": "timestamp", + "DATE": "date", + "DATE32": "date", + } + + if type_name.startswith("LIST<") or type_name.startswith("ARRAY<"): + return "list" + if type_name.startswith("STRUCT<") or type_name.startswith("MAP<"): + return "string" + + return type_map.get(type_name, "string") + + +def _build_iceberg_schema_from_fields( + entity_columns: Sequence, + feature_columns: Sequence, + timestamp_field: Optional[str] = None, +) -> "Schema": + """Build a ``pyiceberg.schema.Schema`` from Feast column descriptors. + + Entity columns become non‑nullable; feature and timestamp columns + are nullable. + """ + from pyiceberg.schema import Schema + from pyiceberg.types import ( + BooleanType, + DateType, + DoubleType, + FloatType, + IntegerType, + LongType, + NestedField, + StringType, + TimestampType, + TimestamptzType, + ) + + PYICEBERG_TYPE_MAP: Dict[str, Any] = { + "boolean": BooleanType(), + "int": IntegerType(), + "long": LongType(), + "float": FloatType(), + "double": DoubleType(), + "string": StringType(), + "binary": StringType(), + "timestamp": TimestampType(), + "timestamptz": TimestamptzType(), + "date": DateType(), + } + + fields: List[NestedField] = [] + seen: set = set() + field_id = 1 + + entity_types: Dict[str, str] = {} + feat_types: Dict[str, str] = {} + + for col in entity_columns: + t = _feast_to_iceberg_type(getattr(col, "dtype", None)) + entity_types[col.name] = t + + for col in feature_columns: + t = _feast_to_iceberg_type(getattr(col, "dtype", None)) + feat_types[col.name] = t + + for col in entity_columns: + if col.name not in seen: + t = entity_types.get(col.name, "string") + ice_type = PYICEBERG_TYPE_MAP.get(t, StringType()) + fields.append(NestedField(field_id, col.name, ice_type, required=True)) + seen.add(col.name) + field_id += 1 + + for col in feature_columns: + if col.name not in seen: + t = feat_types.get(col.name, "string") + ice_type = PYICEBERG_TYPE_MAP.get(t, StringType()) + fields.append(NestedField(field_id, col.name, ice_type, required=False)) + seen.add(col.name) + field_id += 1 + + if timestamp_field and timestamp_field not in seen: + fields.append( + NestedField(field_id, timestamp_field, TimestampType(), required=False) + ) + + return Schema(*fields) + + +def _build_iceberg_properties( + primary_keys: List[str], + timestamp_key: Optional[str], + description: str, + owner: str, + project: str, + tags: Dict[str, str], +) -> Dict[str, str]: + """Build Iceberg table properties from Feast metadata.""" + props: Dict[str, str] = {} + + if primary_keys: + props[FEAST_PRIMARY_KEYS_PROP] = ",".join(primary_keys) + if timestamp_key: + props[FEAST_TIMESTAMP_KEY_PROP] = timestamp_key + if description: + props[FEAST_DESCRIPTION_PROP] = description + if owner: + props[FEAST_OWNER_PROP] = owner + if project: + props[FEAST_PROJECT_PROP] = project + props[FEAST_MANAGED_PROP] = "true" + + for key, value in tags.items(): + safe_key = key.replace(".", "_") + props[f"{_FEAST_TAG_PREFIX}{safe_key}"] = value + + return props + + +def create_feature_table( + catalog: "Catalog", + namespace: str, + table_name: str, + entity_columns: Sequence, + feature_columns: Sequence, + primary_keys: List[str], + timestamp_field: Optional[str], + description: str, + owner: str, + project: str, + tags: Dict[str, str], +) -> None: + """Create an Iceberg table in the catalog with Feast metadata. + + The table schema is derived from Feast entity + feature columns, + and Feast metadata (primary keys, tags, etc.) is stored as Iceberg + table properties. + """ + from pyiceberg.exceptions import NamespaceAlreadyExistsError + + if not namespace_exists(catalog, namespace): + try: + catalog.create_namespace(namespace) + except NamespaceAlreadyExistsError: + pass + + schema = _build_iceberg_schema_from_fields( + entity_columns, feature_columns, timestamp_field + ) + properties = _build_iceberg_properties( + primary_keys, + timestamp_field, + description, + owner, + project, + tags, + ) + + catalog.create_table( + identifier=(namespace, table_name), + schema=schema, + properties=properties, + ) + + +def update_table_properties( + catalog: "Catalog", + namespace: str, + table_name: str, + description: str, + owner: str, + project: str, + tags: Dict[str, str], +) -> None: + """Update Iceberg table properties with current Feast metadata.""" + table = catalog.load_table((namespace, table_name)) + + new_props = _build_iceberg_properties( + primary_keys=[], + timestamp_key=None, + description=description, + owner=owner, + project=project, + tags=tags, + ) + + table.set_properties(**new_props) + + +def write_materialized_data( + catalog: "Catalog", + namespace: str, + table_name: str, + spark_df: SparkDataFrame, + mode: str = "merge", +) -> None: + """Write materialized feature data into an Iceberg table. + + Uses Spark DataFrame writer with Iceberg format for the actual + write. The catalog provides the table metadata; the write is + performed via Spark's native Iceberg integration. + + Args: + catalog: The Iceberg catalog. + namespace: Target namespace (schema). + table_name: Target table name. + spark_df: Spark DataFrame containing the materialized data. + mode: Write mode — ``"merge"``, ``"append"``, or + ``"overwrite"``. Defaults to ``"merge"``. + """ + full_name = f"`{namespace}`.`{table_name}`" + + if mode == "merge": + _merge_into(catalog, namespace, table_name, spark_df) + elif mode == "append": + spark_df.write.format("iceberg").mode("append").save(full_name) + elif mode == "overwrite": + spark_df.write.format("iceberg").mode("overwrite").save(full_name) + else: + raise ValueError(f"Unsupported write mode: {mode}") + + +def _merge_into( + catalog: "Catalog", + namespace: str, + table_name: str, + spark_df: SparkDataFrame, +) -> None: + """Upsert data using Spark MERGE INTO (Delta/Iceberg SQL syntax). + + Requires Spark 3.x with Iceberg support. + """ + full_name = f"{namespace}.{table_name}" + temp_view = f"_feast_merge_{table_name}" + + spark_df.createOrReplaceTempView(temp_view) + + table = catalog.load_table((namespace, table_name)) + primary_keys_str = table.properties.get(FEAST_PRIMARY_KEYS_PROP, "") + pk_cols = [pk.strip() for pk in primary_keys_str.split(",") if pk.strip()] + + if not pk_cols: + spark_df.sql_ctx.sql(f"INSERT INTO {full_name} SELECT * FROM {temp_view}") + return + + join_cond = " AND ".join(f"target.`{pk}` = source.`{pk}`" for pk in pk_cols) + + all_cols = [f.name for f in spark_df.schema.fields] + update_set = ", ".join(f"target.`{c}` = source.`{c}`" for c in all_cols) + insert_cols = ", ".join(f"`{c}`" for c in all_cols) + source_cols = ", ".join(f"source.`{c}`" for c in all_cols) + + merge_sql = ( + f"MERGE INTO {full_name} AS target " + f"USING {temp_view} AS source " + f"ON {join_cond} " + f"WHEN MATCHED THEN UPDATE SET {update_set} " + f"WHEN NOT MATCHED THEN INSERT ({insert_cols}) VALUES ({source_cols})" + ) + + spark_df.sql_ctx.sql(merge_sql) diff --git a/sdk/python/feast/infra/offline_stores/iceberg/iceberg_source.py b/sdk/python/feast/infra/offline_stores/iceberg/iceberg_source.py new file mode 100644 index 00000000000..216da9dd5a2 --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/iceberg/iceberg_source.py @@ -0,0 +1,352 @@ +import json +import logging +from typing import Callable, Dict, Iterable, Optional, Tuple + +from feast.data_source import DataSource +from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto +from feast.repo_config import RepoConfig +from feast.value_type import ValueType + +logger = logging.getLogger(__name__) + + +class IcebergRestCatalogSource(DataSource): + """A DataSource backed by a table in any Iceberg REST Catalog. + + Reading is engine‑agnostic — the catalog metadata (schema, file + listing, credentials) is resolved via ``pyiceberg``, and each + offline store engine (Spark, DuckDB, Trino, …) reads the actual + data files using its native Iceberg support. + """ + + def __init__( + self, + *, + name: str, + endpoint: str, + warehouse: str, + namespace: str, + table: str, + timestamp_field: Optional[str] = None, + token_env_var: str = "ICEBERG_REST_TOKEN", + credential_vending: bool = True, + created_timestamp_column: Optional[str] = None, + field_mapping: Optional[Dict[str, str]] = None, + description: str = "", + tags: Optional[Dict[str, str]] = None, + owner: str = "", + date_partition_column: Optional[str] = None, + date_partition_column_format: str = "%Y-%m-%d", + ): + super().__init__( + name=name, + timestamp_field=timestamp_field, + created_timestamp_column=created_timestamp_column, + field_mapping=field_mapping, + description=description, + date_partition_column=date_partition_column, + tags=tags, + owner=owner, + ) + self.endpoint = endpoint + self.warehouse = warehouse + self.namespace = namespace + self._table = table + self.token_env_var = token_env_var + self.credential_vending = credential_vending + self._date_partition_column_format = date_partition_column_format + + @property + def date_partition_column_format(self): + return self._date_partition_column_format + + @property + def path(self): + return None + + @property + def file_format(self): + return None + + @property + def query(self): + return None + + @property + def full_table_name(self) -> str: + return f"`{self.warehouse}`.`{self.namespace}`.`{self._table}`" + + @staticmethod + def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]: + from feast.type_map import spark_to_feast_value_type + + return spark_to_feast_value_type + + def get_table_query_string(self) -> str: + return self.full_table_name + + def source_type(self) -> DataSourceProto.SourceType.ValueType: + return DataSourceProto.CUSTOM_SOURCE + + def get_table_column_names_and_types( + self, config: RepoConfig + ) -> Iterable[Tuple[str, str]]: + + catalog = load_catalog_from_source(self) + table_ref = catalog.load_table((self.namespace, self._table)) + schema = table_ref.schema() + return ((field.name, str(field.field_type)) for field in schema.fields) + + def validate(self, config: RepoConfig): + self.get_table_column_names_and_types(config) + + @staticmethod + def from_proto(data_source: DataSourceProto) -> "IcebergRestCatalogSource": + assert data_source.HasField("custom_options") + config = json.loads(data_source.custom_options.configuration.decode("utf-8")) + + return IcebergRestCatalogSource( + name=data_source.name, + endpoint=config["endpoint"], + warehouse=config["warehouse"], + namespace=config["namespace"], + table=config["table"], + timestamp_field=data_source.timestamp_field or None, + token_env_var=config.get("token_env_var", "ICEBERG_REST_TOKEN"), + credential_vending=config.get("credential_vending", True), + created_timestamp_column=data_source.created_timestamp_column or None, + field_mapping=dict(data_source.field_mapping) or None, + description=data_source.description or "", + tags=dict(data_source.tags) or None, + owner=data_source.owner or "", + date_partition_column=data_source.date_partition_column or None, + date_partition_column_format=config.get( + "date_partition_column_format", "%Y-%m-%d" + ), + ) + + def _to_proto_impl(self) -> DataSourceProto: + from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto + + config = { + "endpoint": self.endpoint, + "warehouse": self.warehouse, + "namespace": self.namespace, + "table": self._table, + "token_env_var": self.token_env_var, + "credential_vending": self.credential_vending, + "date_partition_column_format": self._date_partition_column_format, + } + + proto = DataSourceProto( + name=self.name, + type=DataSourceProto.CUSTOM_SOURCE, + data_source_class_type="feast.infra.offline_stores.iceberg.iceberg_source.IcebergRestCatalogSource", + field_mapping=self.field_mapping, + date_partition_column=self.date_partition_column or "", + description=self.description or "", + tags=self.tags, + owner=self.owner or "", + custom_options=DataSourceProto.CustomSourceOptions( + configuration=json.dumps(config).encode("utf-8"), + ), + ) + + if self.timestamp_field: + proto.timestamp_field = self.timestamp_field + if self.created_timestamp_column: + proto.created_timestamp_column = self.created_timestamp_column + + return proto + + def __eq__(self, other): + base_eq = super().__eq__(other) + if not base_eq: + return False + if not isinstance(other, IcebergRestCatalogSource): + return False + return ( + self.endpoint == other.endpoint + and self.warehouse == other.warehouse + and self.namespace == other.namespace + and self._table == other._table + and self.token_env_var == other.token_env_var + and self.credential_vending == other.credential_vending + and self._date_partition_column_format + == other._date_partition_column_format + ) + + def __hash__(self): + return super().__hash__() + + +class UnityCatalogSource(IcebergRestCatalogSource): + """An Iceberg REST Catalog source specialised for Databricks Unity Catalog. + + Adds UC‑specific governance fields on top of the base + ``IcebergRestCatalogSource``: + + * ``register_as_feature_table`` — whether to register this + FeatureView as a UC feature table during ``feast apply``. + * ``sync_lineage`` — whether to emit OpenLineage events with UC + fully‑qualified table references during materialization. + """ + + def __init__( + self, + *, + name: str, + endpoint: str, + warehouse: str, + namespace: str, + table: str, + timestamp_field: Optional[str] = None, + token_env_var: str = "DATABRICKS_TOKEN", + credential_vending: bool = True, + cluster_id: Optional[str] = None, + register_as_feature_table: bool = True, + sync_lineage: bool = True, + created_timestamp_column: Optional[str] = None, + field_mapping: Optional[Dict[str, str]] = None, + description: str = "", + tags: Optional[Dict[str, str]] = None, + owner: str = "", + date_partition_column: Optional[str] = None, + date_partition_column_format: str = "%Y-%m-%d", + ): + super().__init__( + name=name, + endpoint=endpoint, + warehouse=warehouse, + namespace=namespace, + table=table, + timestamp_field=timestamp_field, + token_env_var=token_env_var, + credential_vending=credential_vending, + created_timestamp_column=created_timestamp_column, + field_mapping=field_mapping, + description=description, + tags=tags, + owner=owner, + date_partition_column=date_partition_column, + date_partition_column_format=date_partition_column_format, + ) + self.cluster_id = cluster_id + self.register_as_feature_table = register_as_feature_table + self.sync_lineage = sync_lineage + + def source_type(self) -> DataSourceProto.SourceType.ValueType: + return DataSourceProto.CUSTOM_SOURCE + + @staticmethod + def from_proto(data_source: DataSourceProto) -> "UnityCatalogSource": + assert data_source.HasField("custom_options") + config = json.loads(data_source.custom_options.configuration.decode("utf-8")) + + return UnityCatalogSource( + name=data_source.name, + endpoint=config["endpoint"], + warehouse=config["warehouse"], + namespace=config["namespace"], + table=config["table"], + timestamp_field=data_source.timestamp_field or None, + token_env_var=config.get("token_env_var", "DATABRICKS_TOKEN"), + credential_vending=config.get("credential_vending", True), + cluster_id=config.get("cluster_id"), + register_as_feature_table=config.get("register_as_feature_table", True), + sync_lineage=config.get("sync_lineage", True), + created_timestamp_column=data_source.created_timestamp_column or None, + field_mapping=dict(data_source.field_mapping) or None, + description=data_source.description or "", + tags=dict(data_source.tags) or None, + owner=data_source.owner or "", + date_partition_column=data_source.date_partition_column or None, + date_partition_column_format=config.get( + "date_partition_column_format", "%Y-%m-%d" + ), + ) + + def _to_proto_impl(self) -> DataSourceProto: + from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto + + config = { + "endpoint": self.endpoint, + "warehouse": self.warehouse, + "namespace": self.namespace, + "table": self._table, + "token_env_var": self.token_env_var, + "credential_vending": self.credential_vending, + "cluster_id": self.cluster_id, + "register_as_feature_table": self.register_as_feature_table, + "sync_lineage": self.sync_lineage, + "date_partition_column_format": self._date_partition_column_format, + } + + proto = DataSourceProto( + name=self.name, + type=DataSourceProto.CUSTOM_SOURCE, + data_source_class_type="feast.infra.offline_stores.iceberg.iceberg_source.UnityCatalogSource", + field_mapping=self.field_mapping, + date_partition_column=self.date_partition_column or "", + description=self.description or "", + tags=self.tags, + owner=self.owner or "", + custom_options=DataSourceProto.CustomSourceOptions( + configuration=json.dumps(config).encode("utf-8"), + ), + ) + + if self.timestamp_field: + proto.timestamp_field = self.timestamp_field + if self.created_timestamp_column: + proto.created_timestamp_column = self.created_timestamp_column + + return proto + + def __eq__(self, other): + base_eq = super().__eq__(other) + if not base_eq: + return False + if not isinstance(other, UnityCatalogSource): + return False + return ( + self.cluster_id == other.cluster_id + and self.register_as_feature_table == other.register_as_feature_table + and self.sync_lineage == other.sync_lineage + ) + + def __hash__(self): + return super().__hash__() + + +def load_catalog_from_source( + source: IcebergRestCatalogSource, +): + """Build a ``pyiceberg.catalog.Catalog`` from an Iceberg source. + + This helper resolves the token from the environment, constructs + the catalog properties, and delegates to ``pyiceberg.catalog.load_catalog()``. + """ + import os + + from pyiceberg.catalog import load_catalog + + token = os.environ.get(source.token_env_var, "") + + catalog_type = "rest" + if isinstance(source, UnityCatalogSource): + catalog_type = "databricks" + + props: Dict[str, str] = { + "uri": source.endpoint, + "warehouse": source.warehouse, + "token": token, + } + if source.credential_vending: + props["credential-vending"] = "true" + + return load_catalog( + name=source.warehouse, + type=catalog_type, + **props, + ) diff --git a/sdk/python/feast/infra/offline_stores/iceberg/registration.py b/sdk/python/feast/infra/offline_stores/iceberg/registration.py new file mode 100644 index 00000000000..72b7c2927e4 --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/iceberg/registration.py @@ -0,0 +1,188 @@ +"""Iceberg REST Catalog registration for ``feast apply``. + +When the offline store is configured with an Iceberg REST Catalog, this +module registers (or updates) each ``FeatureView`` as a table in the +catalog. FeatureView entities become primary keys (stored as an Iceberg +table property), and tags/description/owner are synced to Iceberg table +properties. + +Per‑FeatureView opt‑out and overrides are controlled via ``FeatureView`` +``tags``: + +* ``uc.register_as_feature_table`` — ``"false"`` skips a specific view. +* ``uc.catalog`` / ``uc.schema`` / ``uc.table`` — override the UC path. + +Global defaults come from ``UCRegistrationConfig`` (inside +``DatabricksUCOfflineStoreConfig``). +""" + +import logging +from typing import Dict, List, Optional, Tuple + +import click + +from feast import FeatureView +from feast.infra.offline_stores.iceberg.catalog_config import IcebergCatalogConfig +from feast.infra.offline_stores.iceberg.catalog_manager import ( + create_feature_table, + load_catalog, + table_exists, + update_table_properties, +) + +logger = logging.getLogger(__name__) + +# FeatureView tag keys for per‑FV configuration +_REGISTER_AS_FEATURE_TABLE_KEY = "uc.register_as_feature_table" +_CATALOG_KEY = "uc.catalog" +_SCHEMA_KEY = "uc.schema" +_TABLE_KEY = "uc.table" + + +def _should_register(fv: FeatureView) -> bool: + """Return ``True`` unless the feature view opts out via tags.""" + return fv.tags.get(_REGISTER_AS_FEATURE_TABLE_KEY, "true").lower() != "false" + + +def _resolve_uc_path( + fv: FeatureView, + default_catalog: Optional[str], + default_schema: Optional[str], +) -> Tuple[Optional[str], Optional[str], str]: + """Resolve the (catalog, schema, table_name) for a feature view. + + Prioritises per‑FV tag overrides, then global defaults, then the + feature view name as the table name. + """ + catalog = fv.tags.get(_CATALOG_KEY) or default_catalog + schema = fv.tags.get(_SCHEMA_KEY) or default_schema + table = fv.tags.get(_TABLE_KEY) or fv.name + + table = table.replace("-", "_").replace(".", "_").replace(" ", "_") + return catalog, schema, table + + +def _get_primary_keys(fv: FeatureView) -> List[str]: + """Extract primary key column names from entity columns.""" + return [col.name for col in fv.entity_columns] + + +def _build_tags(fv: FeatureView, project: str) -> Dict[str, str]: + """Build table tags, excluding internal ``uc.*`` keys.""" + tags: Dict[str, str] = {} + for key, value in fv.tags.items(): + if not key.startswith("uc."): + tags[key] = value + tags["feast_managed"] = "feast" + tags["feast_project"] = project + return tags + + +def _register_single_feature_view( + catalog, + fv: FeatureView, + project: str, + default_catalog: Optional[str], + default_schema: Optional[str], +) -> None: + """Register or update a single FeatureView as a catalog table.""" + if not _should_register(fv): + click.echo(f" ⊘ Skipping catalog registration for {fv.name} (opt‑out)") + return + + catalog_name, schema, table = _resolve_uc_path(fv, default_catalog, default_schema) + if not catalog_name or not schema: + click.echo( + f" ⚠ Cannot register {fv.name}: missing catalog or schema. " + f"Set default_catalog/default_schema in feature_store.yaml or " + f"use tags uc.catalog/uc.schema." + ) + return + + namespace = schema + table_name = table + primary_keys = _get_primary_keys(fv) + timestamp_field = getattr(fv.batch_source, "timestamp_field", None) + description = fv.description or "" + owner = fv.owner or "" + tags = _build_tags(fv, project) + + if table_exists(catalog, namespace, table_name): + update_table_properties( + catalog=catalog, + namespace=namespace, + table_name=table_name, + description=description, + owner=owner, + project=project, + tags=tags, + ) + click.echo( + f" ✓ Updated catalog table: {catalog_name}.{namespace}.{table_name}" + ) + else: + create_feature_table( + catalog=catalog, + namespace=namespace, + table_name=table_name, + entity_columns=fv.entity_columns, + feature_columns=fv.features, + primary_keys=primary_keys, + timestamp_field=timestamp_field, + description=description, + owner=owner, + project=project, + tags=tags, + ) + click.echo( + f" ✓ Created catalog table: {catalog_name}.{namespace}.{table_name}" + ) + + +def register_uc_feature_tables( + catalog_config: IcebergCatalogConfig, + feature_views: List[FeatureView], + project: str, +) -> None: + """Register or update FeatureViews as Iceberg REST Catalog tables. + + Skips silently when: + - ``pyiceberg`` is not installed. + - The catalog cannot be reached. + - All feature views opt out. + + Per‑feature‑view errors are logged and do not halt the apply. + """ + try: + import pyiceberg # noqa: F401 + except ImportError: + logger.info( + "pyiceberg is not installed; skipping catalog registration. " + "Install with: pip install pyiceberg[pyarrow]" + ) + return + + try: + catalog = load_catalog(catalog_config) + except Exception as e: + logger.warning("Could not create Iceberg catalog for registration: %s", e) + return + + for fv in feature_views: + try: + _register_single_feature_view( + catalog, + fv, + project, + catalog_config.warehouse, + catalog_config.namespace, + ) + except Exception: + logger.exception( + "Failed to register catalog table for FeatureView '%s'", + fv.name, + ) + click.echo( + f" ✗ Failed to register catalog table: {fv.name} " + f"(check logs for details)" + ) diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py index 07bfd1d5416..73a6506c7cc 100644 --- a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py @@ -12,33 +12,70 @@ from feast.infra.offline_stores.contrib.spark_offline_store.spark_source import ( SparkSource, ) +from feast.infra.offline_stores.iceberg.catalog_config import IcebergCatalogConfig from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig from feast.repo_config import RepoConfig +def _catalog_config(**overrides) -> IcebergCatalogConfig: + defaults = { + "type": "rest", + "endpoint": "https://adb-12345.azuredatabricks.net/api/2.1/unity-catalog/iceberg", + "warehouse": "main", + "namespace": "default", + } + return IcebergCatalogConfig(**{**defaults, **overrides}) + + def test_config_parsing(): config_dict = { "type": "databricks_uc", - "workspace_host": "adb-12345.azuredatabricks.net", - "token": "dapi123456", - "cluster_id": "0123-4567-abcde", - "default_catalog": "main", - "default_schema": "default", + "catalog": { + "type": "rest", + "endpoint": "https://adb-12345.azuredatabricks.net/api/2.1/unity-catalog/iceberg", + "warehouse": "main", + "namespace": "default", + "token_env_var": "DATABRICKS_TOKEN", + }, "spark_conf": {"spark.sql.shuffle.partitions": "10"}, } config = DatabricksUCOfflineStoreConfig(**config_dict) assert config.type == "databricks_uc" - assert config.workspace_host == "adb-12345.azuredatabricks.net" - assert config.token == "dapi123456" - assert config.cluster_id == "0123-4567-abcde" - assert config.default_catalog == "main" - assert config.default_schema == "default" + assert config.catalog.warehouse == "main" + assert config.catalog.namespace == "default" + assert config.catalog.token_env_var == "DATABRICKS_TOKEN" assert config.spark_conf == {"spark.sql.shuffle.partitions": "10"} +def test_config_with_uc_registration(): + config_dict = { + "type": "databricks_uc", + "catalog": { + "type": "rest", + "endpoint": "https://adb-12345.azuredatabricks.net/api/2.1/unity-catalog/iceberg", + "warehouse": "prod_ml", + "namespace": "features", + }, + "uc_registration": { + "enabled": True, + "catalog": "override_cat", + "uc_schema": "override_sch", + }, + } + config = DatabricksUCOfflineStoreConfig(**config_dict) + assert config.uc_registration is not None + assert config.uc_registration.enabled is True + assert config.uc_registration.catalog == "override_cat" + assert config.uc_registration.uc_schema == "override_sch" + + def test_config_forbidden_extra(): with pytest.raises(ValidationError): - DatabricksUCOfflineStoreConfig(type="databricks_uc", invalid_key="some_val") + DatabricksUCOfflineStoreConfig( + type="databricks_uc", + catalog=_catalog_config(), + invalid_key="some_val", + ) @patch("pyspark.sql.SparkSession.getActiveSession") @@ -48,8 +85,7 @@ def test_get_databricks_session_active(mock_get_active): config = DatabricksUCOfflineStoreConfig( type="databricks_uc", - default_catalog="my_catalog", - default_schema="my_schema", + catalog=_catalog_config(warehouse="my_catalog", namespace="my_schema"), ) session = get_databricks_session(config) @@ -73,9 +109,11 @@ def test_get_databricks_session_new_remote(mock_builder, mock_get_active): config = DatabricksUCOfflineStoreConfig( type="databricks_uc", - workspace_host="https://adb-12345.azuredatabricks.net", - token="dapi123", - cluster_id="0123-4567-abcde", + catalog=_catalog_config( + endpoint="https://adb-12345.azuredatabricks.net/api/2.1/unity-catalog/iceberg", + warehouse="main", + namespace="default", + ), spark_conf={"spark.some.option": "value"}, ) @@ -83,7 +121,7 @@ def test_get_databricks_session_new_remote(mock_builder, mock_get_active): assert session == mock_session mock_builder.remote.assert_called_once_with( - "sc://adb-12345.azuredatabricks.net:443/;token=dapi123;x-databricks-cluster-id=0123-4567-abcde" + "sc://adb-12345.azuredatabricks.net:443/" ) @@ -104,8 +142,9 @@ def test_get_historical_features_delegation(mock_parent_features, mock_get_sessi online_store=SqliteOnlineStoreConfig(type="sqlite"), offline_store=DatabricksUCOfflineStoreConfig( type="databricks_uc", - workspace_host="adb-123.databricks.com", - cluster_id="123", + catalog=_catalog_config( + endpoint="https://adb-123.databricks.com/api/2.1/unity-catalog/iceberg", + ), ), ) @@ -154,8 +193,9 @@ def test_pull_latest_from_table_or_query_delegation( online_store=SqliteOnlineStoreConfig(type="sqlite"), offline_store=DatabricksUCOfflineStoreConfig( type="databricks_uc", - workspace_host="adb-123.databricks.com", - cluster_id="123", + catalog=_catalog_config( + endpoint="https://adb-123.databricks.com/api/2.1/unity-catalog/iceberg", + ), ), ) diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_iceberg_source.py b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_iceberg_source.py new file mode 100644 index 00000000000..c844f3330af --- /dev/null +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_iceberg_source.py @@ -0,0 +1,210 @@ +from unittest.mock import MagicMock, patch + +from feast.infra.offline_stores.iceberg.iceberg_source import ( + IcebergRestCatalogSource, + UnityCatalogSource, + load_catalog_from_source, +) +from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto + + +def _make_iceberg_source(**overrides): + defaults = { + "name": "test_source", + "endpoint": "https://test.databricks.net/api/2.1/unity-catalog/iceberg", + "warehouse": "test_cat", + "namespace": "test_sch", + "table": "test_table", + "timestamp_field": "event_ts", + } + return IcebergRestCatalogSource(**{**defaults, **overrides}) + + +def test_iceberg_source_init(): + src = _make_iceberg_source() + assert src.name == "test_source" + assert src.endpoint == "https://test.databricks.net/api/2.1/unity-catalog/iceberg" + assert src.warehouse == "test_cat" + assert src.namespace == "test_sch" + assert src._table == "test_table" + assert src.timestamp_field == "event_ts" + assert src.token_env_var == "ICEBERG_REST_TOKEN" + assert src.credential_vending is True + + +def test_iceberg_source_full_table_name(): + src = _make_iceberg_source() + assert src.full_table_name == "`test_cat`.`test_sch`.`test_table`" + + +def test_iceberg_source_get_table_query_string(): + src = _make_iceberg_source() + assert src.get_table_query_string() == "`test_cat`.`test_sch`.`test_table`" + + +def test_iceberg_source_source_type(): + src = _make_iceberg_source() + assert src.source_type() == DataSourceProto.CUSTOM_SOURCE + + +def test_iceberg_source_proto_roundtrip(): + src = _make_iceberg_source( + name="proto_test", + timestamp_field="ts", + description="a test source", + owner="user@test.com", + tags={"env": "test"}, + ) + proto = src.to_proto() + assert proto.name == "proto_test" + assert proto.type == DataSourceProto.CUSTOM_SOURCE + assert proto.owner == "user@test.com" + assert proto.description == "a test source" + + restored = IcebergRestCatalogSource.from_proto(proto) + assert restored == src + assert restored.name == src.name + assert restored.endpoint == src.endpoint + assert restored.warehouse == src.warehouse + assert restored.namespace == src.namespace + assert restored._table == src._table + + +def test_iceberg_source_equality(): + src1 = _make_iceberg_source(name="src1", table="t1") + src2 = _make_iceberg_source(name="src1", table="t1") + assert src1 == src2 + + src3 = _make_iceberg_source(name="diff", table="t2") + assert src1 != src3 + + +def test_iceberg_source_path_file_format_query(): + src = _make_iceberg_source() + assert src.path is None + assert src.file_format is None + assert src.query is None + + +def test_iceberg_source_proto_class_type(): + src = _make_iceberg_source() + proto = src.to_proto() + assert ( + proto.data_source_class_type + == "feast.infra.offline_stores.iceberg.iceberg_source.IcebergRestCatalogSource" + ) + + +def test_unity_catalog_source_init(): + src = UnityCatalogSource( + name="uc_source", + endpoint="https://test.databricks.net/api/2.1/unity-catalog/iceberg", + warehouse="uc_cat", + namespace="uc_sch", + table="uc_table", + timestamp_field="ts", + register_as_feature_table=True, + sync_lineage=True, + ) + assert src.register_as_feature_table is True + assert src.sync_lineage is True + assert src.token_env_var == "DATABRICKS_TOKEN" + + +def test_unity_catalog_source_proto_roundtrip(): + src = UnityCatalogSource( + name="uc_test", + endpoint="https://test.databricks.net/api/2.1/unity-catalog/iceberg", + warehouse="uc_cat", + namespace="uc_sch", + table="uc_table", + timestamp_field="ts", + cluster_id="0123-4567-abcde", + register_as_feature_table=True, + sync_lineage=True, + ) + proto = src.to_proto() + restored = UnityCatalogSource.from_proto(proto) + assert restored == src + assert restored.cluster_id == "0123-4567-abcde" + assert restored.register_as_feature_table is True + assert restored.sync_lineage is True + + +def test_unity_catalog_source_equality(): + src1 = UnityCatalogSource( + name="u1", + endpoint="https://test.databricks.net", + warehouse="cat", + namespace="sch", + table="t1", + cluster_id="c1", + ) + src2 = UnityCatalogSource( + name="u1", + endpoint="https://test.databricks.net", + warehouse="cat", + namespace="sch", + table="t1", + cluster_id="c1", + ) + src3 = UnityCatalogSource( + name="u1", + endpoint="https://test.databricks.net", + warehouse="cat", + namespace="sch", + table="t1", + cluster_id="c2", + ) + assert src1 == src2 + assert src1 != src3 + + +def test_unity_catalog_source_proto_class_type(): + src = UnityCatalogSource( + name="uc_test", + endpoint="https://test.databricks.net/api/2.1/unity-catalog/iceberg", + warehouse="cat", + namespace="sch", + table="t", + ) + proto = src.to_proto() + assert ( + proto.data_source_class_type + == "feast.infra.offline_stores.iceberg.iceberg_source.UnityCatalogSource" + ) + + +@patch("feast.infra.offline_stores.iceberg.iceberg_source.load_catalog") +def test_load_catalog_from_source_iceberg(mock_load_catalog): + mock_catalog = MagicMock() + mock_load_catalog.return_value = mock_catalog + + src = _make_iceberg_source() + catalog = load_catalog_from_source(src) + + assert catalog == mock_catalog + mock_load_catalog.assert_called_once() + call_kwargs = mock_load_catalog.call_args[1] + assert call_kwargs["type"] == "rest" + assert src.endpoint in call_kwargs.get("uri", "") + + +@patch("feast.infra.offline_stores.iceberg.iceberg_source.load_catalog") +def test_load_catalog_from_source_unity(mock_load_catalog): + mock_catalog = MagicMock() + mock_load_catalog.return_value = mock_catalog + + src = UnityCatalogSource( + name="uc", + endpoint="https://test.databricks.net/api/2.1/unity-catalog/iceberg", + warehouse="cat", + namespace="sch", + table="t", + ) + catalog = load_catalog_from_source(src) + + assert catalog == mock_catalog + mock_load_catalog.assert_called_once() + call_kwargs = mock_load_catalog.call_args[1] + assert call_kwargs["type"] == "databricks" diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py index c394b1c1353..cd0ebe6adee 100644 --- a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py @@ -1,28 +1,10 @@ from unittest.mock import MagicMock, patch -from pyspark.sql.types import ( - ArrayType, - BinaryType, - BooleanType, - DoubleType, - FloatType, - IntegerType, - LongType, - StringType, - StructType, - TimestampType, -) - from feast import FeatureView -from feast.infra.offline_stores.contrib.spark_offline_store.databricks_uc import ( - DatabricksUCOfflineStoreConfig, - UCRegistrationConfig, -) -from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( - _build_full_table_name, - _build_spark_schema, - _build_uc_tags, - _feast_to_spark_type_simple, +from feast.infra.offline_stores.iceberg.catalog_config import IcebergCatalogConfig +from feast.infra.offline_stores.iceberg.catalog_manager import _feast_to_iceberg_type +from feast.infra.offline_stores.iceberg.registration import ( + _build_tags, _get_primary_keys, _resolve_uc_path, _should_register, @@ -33,8 +15,6 @@ def make_mock_field(name: str, dtype_str: str = "", nullable: bool = True): field = MagicMock() field.name = name - # _feast_to_spark_type_simple does str(dtype).upper() - # When dtype_str is empty, dtype will be None -> StringType fallback if dtype_str: class FakeDtype: @@ -47,76 +27,56 @@ def __str__(self): return field -def test_feast_to_spark_type_simple_none(): - field = make_mock_field("col", "") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, StringType) +def test_feast_to_iceberg_type_none(): + assert _feast_to_iceberg_type(None) == "string" -def test_feast_to_spark_type_simple_int32(): - field = make_mock_field("col", "Int32") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, IntegerType) +def test_feast_to_iceberg_type_int32(): + assert _feast_to_iceberg_type(make_mock_field("col", "Int32").dtype) == "int" -def test_feast_to_spark_type_simple_int64(): - field = make_mock_field("col", "Int64") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, LongType) +def test_feast_to_iceberg_type_int64(): + assert _feast_to_iceberg_type(make_mock_field("col", "Int64").dtype) == "long" -def test_feast_to_spark_type_simple_float32(): - field = make_mock_field("col", "Float32") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, FloatType) +def test_feast_to_iceberg_type_float32(): + assert _feast_to_iceberg_type(make_mock_field("col", "Float32").dtype) == "float" -def test_feast_to_spark_type_simple_float64(): - field = make_mock_field("col", "Float64") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, DoubleType) +def test_feast_to_iceberg_type_float64(): + assert _feast_to_iceberg_type(make_mock_field("col", "Float64").dtype) == "double" -def test_feast_to_spark_type_simple_string(): - field = make_mock_field("col", "String") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, StringType) +def test_feast_to_iceberg_type_string(): + assert _feast_to_iceberg_type(make_mock_field("col", "String").dtype) == "string" -def test_feast_to_spark_type_simple_bool(): - field = make_mock_field("col", "Bool") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, BooleanType) +def test_feast_to_iceberg_type_bool(): + assert _feast_to_iceberg_type(make_mock_field("col", "Bool").dtype) == "boolean" -def test_feast_to_spark_type_simple_bytes(): - field = make_mock_field("col", "Bytes") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, BinaryType) +def test_feast_to_iceberg_type_bytes(): + assert _feast_to_iceberg_type(make_mock_field("col", "Bytes").dtype) == "binary" -def test_feast_to_spark_type_simple_unix_timestamp(): - field = make_mock_field("col", "UnixTimestamp") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, TimestampType) +def test_feast_to_iceberg_type_unix_timestamp(): + t = _feast_to_iceberg_type(make_mock_field("col", "UnixTimestamp").dtype) + assert t == "timestamp" -def test_feast_to_spark_type_simple_list(): - field = make_mock_field("col", "List") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, ArrayType) +def test_feast_to_iceberg_type_list(): + t = _feast_to_iceberg_type(make_mock_field("col", "List").dtype) + assert t == "list" -def test_feast_to_spark_type_simple_array(): - field = make_mock_field("col", "Array") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, ArrayType) +def test_feast_to_iceberg_type_array(): + t = _feast_to_iceberg_type(make_mock_field("col", "Array").dtype) + assert t == "list" -def test_feast_to_spark_type_simple_unknown(): - field = make_mock_field("col", "UnknownType") - result = _feast_to_spark_type_simple(field) - assert isinstance(result, StringType) +def test_feast_to_iceberg_type_unknown(): + t = _feast_to_iceberg_type(make_mock_field("col", "UnknownType").dtype) + assert t == "string" def test_should_register_default(): @@ -185,14 +145,14 @@ def test_resolve_uc_path_tag_overrides_none_default(): assert table == "fv_name" -def test_build_uc_tags(): +def test_build_tags(): fv = MagicMock(spec=FeatureView) fv.tags = { "env": "prod", "uc.register_as_feature_table": "false", "owner": "team_a", } - tags = _build_uc_tags(fv, "my_project") + tags = _build_tags(fv, "my_project") assert tags["env"] == "prod" assert tags["owner"] == "team_a" assert "uc.register_as_feature_table" not in tags @@ -200,10 +160,10 @@ def test_build_uc_tags(): assert tags["feast_project"] == "my_project" -def test_build_uc_tags_empty_tags(): +def test_build_tags_empty_tags(): fv = MagicMock(spec=FeatureView) fv.tags = {} - tags = _build_uc_tags(fv, "proj") + tags = _build_tags(fv, "proj") assert tags["feast_managed"] == "feast" assert tags["feast_project"] == "proj" @@ -222,100 +182,64 @@ def test_get_primary_keys_empty(): assert _get_primary_keys(fv) == [] -def test_build_full_table_name_all_parts(): - assert _build_full_table_name("cat", "sch", "table") == "`cat`.`sch`.`table`" - - -def test_build_full_table_name_no_catalog(): - assert _build_full_table_name(None, "sch", "table") == "`sch`.`table`" - - -def test_build_full_table_name_no_schema(): - assert _build_full_table_name("cat", None, "table") == "`cat`.`table`" - - -def test_build_full_table_name_only_table(): - assert _build_full_table_name(None, None, "table") == "`table`" - - -def test_build_spark_schema(): - fv = MagicMock(spec=FeatureView) - id_field = make_mock_field("entity_id", "Int32") - feat_field = make_mock_field("feature_val", "Float64") - fv.entity_columns = [id_field] - fv.features = [feat_field] - fv.batch_source = MagicMock() - fv.batch_source.timestamp_field = "event_ts" +def _make_catalog_config(**overrides) -> IcebergCatalogConfig: + defaults = { + "type": "rest", + "endpoint": "https://test.databricks.net/api/2.1/unity-catalog/iceberg", + "warehouse": "test_cat", + "namespace": "test_sch", + } + merged = {**defaults, **overrides} + return IcebergCatalogConfig(**merged) - schema = _build_spark_schema(fv) - assert isinstance(schema, StructType) - fields = schema.fields - assert fields[0].name == "entity_id" - assert isinstance(fields[0].dataType, IntegerType) - assert fields[0].nullable is False +def test_register_uc_feature_tables_skips_when_missing_pyiceberg(): + config = _make_catalog_config() + with patch.dict("sys.modules", {"pyiceberg": None}): + register_uc_feature_tables(config, [], "test_project") + # No error means skip silently - assert fields[1].name == "feature_val" - assert isinstance(fields[1].dataType, DoubleType) - assert fields[1].nullable is True - assert fields[2].name == "event_ts" - assert isinstance(fields[2].dataType, TimestampType) - assert fields[2].nullable is True +@patch("feast.infra.offline_stores.iceberg.registration.load_catalog") +def test_register_uc_feature_tables_creates_new_table(mock_load_catalog): + mock_catalog = MagicMock() + mock_load_catalog.return_value = mock_catalog + mock_catalog.load_table.side_effect = Exception("table not found") + config = _make_catalog_config() -def test_build_spark_schema_no_timestamp(): fv = MagicMock(spec=FeatureView) - id_field = make_mock_field("id", "Int64") - fv.entity_columns = [id_field] + fv.tags = {} + fv.name = "test_fv" + fv.entity_columns = [] fv.features = [] fv.batch_source = MagicMock() fv.batch_source.timestamp_field = None + fv.description = "test description" + fv.owner = "user1" - schema = _build_spark_schema(fv) - assert len(schema.fields) == 1 - assert schema.fields[0].name == "id" - - -def test_register_uc_feature_tables_skips_when_disabled(): - config = DatabricksUCOfflineStoreConfig(type="databricks_uc") - register_uc_feature_tables(config, [], "test_project") - # No error means success (skip silently) + with patch( + "feast.infra.offline_stores.iceberg.registration.table_exists", + return_value=False, + ): + with patch( + "feast.infra.offline_stores.iceberg.registration.create_feature_table" + ) as mock_create: + register_uc_feature_tables(config, [fv], "proj") + mock_create.assert_called_once() + call_kwargs = mock_create.call_args[1] + assert call_kwargs["namespace"] == "test_sch" + assert call_kwargs["table_name"] == "test_fv" + assert call_kwargs["description"] == "test description" -def test_register_uc_feature_tables_skips_when_no_uc_config(): - config = DatabricksUCOfflineStoreConfig( - type="databricks_uc", - uc_registration=UCRegistrationConfig(enabled=False), - ) - register_uc_feature_tables(config, [], "test_project") - # No error means success +@patch("feast.infra.offline_stores.iceberg.registration.load_catalog") +def test_register_uc_feature_tables_updates_existing(mock_load_catalog): + mock_catalog = MagicMock() + mock_load_catalog.return_value = mock_catalog -def _run_with_mock_fe(config, fvs, project, table_exists: bool = False): - """Helper to run register_uc_feature_tables with mocked external deps.""" - fe_module = MagicMock() - fe_module.FeatureEngineeringClient = MagicMock() - with patch.dict("sys.modules", {"databricks.feature_engineering": fe_module}): - with patch( - "feast.infra.offline_stores.contrib.spark_offline_store.uc_registration.get_databricks_session" - ) as mock_get_session: - fe_client = MagicMock() - fe_module.FeatureEngineeringClient.return_value = fe_client - mock_session = MagicMock() - mock_session.catalog.tableExists.return_value = table_exists - mock_get_session.return_value = mock_session - register_uc_feature_tables(config, fvs, project) - return fe_client, mock_session - - -def test_register_uc_feature_tables_creates_new_table(): - config = DatabricksUCOfflineStoreConfig( - type="databricks_uc", - default_catalog="cat", - default_schema="sch", - uc_registration=UCRegistrationConfig(enabled=True), - ) + config = _make_catalog_config() fv = MagicMock(spec=FeatureView) fv.tags = {} @@ -327,77 +251,57 @@ def test_register_uc_feature_tables_creates_new_table(): fv.description = "test description" fv.owner = "user1" - fe_client, mock_session = _run_with_mock_fe( - config, [fv], "proj", table_exists=False - ) + with patch( + "feast.infra.offline_stores.iceberg.registration.table_exists", + return_value=True, + ): + with patch( + "feast.infra.offline_stores.iceberg.registration.update_table_properties" + ) as mock_update: + register_uc_feature_tables(config, [fv], "proj") - fe_client.create_table.assert_called_once() - call_kwargs = fe_client.create_table.call_args[1] - assert call_kwargs["name"] == "`cat`.`sch`.`test_fv`" - assert call_kwargs["primary_keys"] == [] + mock_update.assert_called_once() + call_kwargs = mock_update.call_args[1] + assert call_kwargs["namespace"] == "test_sch" + assert call_kwargs["table_name"] == "test_fv" assert call_kwargs["description"] == "test description" -@patch( - "feast.infra.offline_stores.contrib.spark_offline_store.uc_registration.get_databricks_session" -) -def test_register_uc_feature_tables_updates_existing( - mock_get_session, -): - mock_session = MagicMock() - mock_get_session.return_value = mock_session - mock_session.catalog.tableExists.return_value = True - - fe_module = MagicMock() - fe_client = MagicMock() - fe_module.FeatureEngineeringClient.return_value = fe_client +@patch("feast.infra.offline_stores.iceberg.registration.load_catalog") +def test_register_uc_feature_tables_skips_opt_out(mock_load_catalog): + mock_catalog = MagicMock() + mock_load_catalog.return_value = mock_catalog - config = DatabricksUCOfflineStoreConfig( - type="databricks_uc", - default_catalog="cat", - default_schema="sch", - uc_registration=UCRegistrationConfig(enabled=True), - ) + config = _make_catalog_config() fv = MagicMock(spec=FeatureView) - fv.tags = {} - fv.name = "test_fv" + fv.tags = {"uc.register_as_feature_table": "false"} + fv.name = "opt_out_fv" fv.entity_columns = [] fv.features = [] fv.batch_source = MagicMock() fv.batch_source.timestamp_field = None - fv.description = "test description" - fv.owner = "user1" + fv.description = "" + fv.owner = None - with patch.dict("sys.modules", {"databricks.feature_engineering": fe_module}): + with patch( + "feast.infra.offline_stores.iceberg.registration.create_feature_table" + ) as mock_create: register_uc_feature_tables(config, [fv], "proj") - fe_client.create_table.assert_not_called() - mock_session.sql.assert_any_call( - "COMMENT ON TABLE `cat`.`sch`.`test_fv` IS 'test description'" - ) + mock_create.assert_not_called() -@patch( - "feast.infra.offline_stores.contrib.spark_offline_store.uc_registration.get_databricks_session" -) -def test_register_uc_feature_tables_skips_opt_out( - mock_get_session, +@patch("feast.infra.offline_stores.iceberg.registration.load_catalog") +def test_register_uc_feature_tables_skips_when_missing_catalog( + mock_load_catalog, ): - mock_session = MagicMock() - mock_get_session.return_value = mock_session - mock_session.catalog.tableExists.return_value = False - - config = DatabricksUCOfflineStoreConfig( - type="databricks_uc", - default_catalog="cat", - default_schema="sch", - uc_registration=UCRegistrationConfig(enabled=True), - ) + mock_catalog = MagicMock() + mock_load_catalog.return_value = mock_catalog fv = MagicMock(spec=FeatureView) - fv.tags = {"uc.register_as_feature_table": "false"} - fv.name = "opt_out_fv" + fv.tags = {} + fv.name = "no_cat_fv" fv.entity_columns = [] fv.features = [] fv.batch_source = MagicMock() @@ -405,44 +309,16 @@ def test_register_uc_feature_tables_skips_opt_out( fv.description = "" fv.owner = None - fe_module = MagicMock() - fe_client = MagicMock() - fe_module.FeatureEngineeringClient.return_value = fe_client - with patch.dict("sys.modules", {"databricks.feature_engineering": fe_module}): - register_uc_feature_tables(config, [fv], "proj") - - # Per-view opt-out means create_table should NOT be called - fe_client.create_table.assert_not_called() + # Use a config with empty warehouse/namespace to simulate missing catalog info + config_no_warehouse = _make_catalog_config(warehouse="", namespace="") - -def test_register_uc_feature_tables_skips_when_missing_catalog(): - fe_module = MagicMock() - fe_client = MagicMock() - fe_module.FeatureEngineeringClient.return_value = fe_client - with patch.dict("sys.modules", {"databricks.feature_engineering": fe_module}): + with patch( + "feast.infra.offline_stores.iceberg.registration.create_feature_table" + ) as mock_create: with patch( - "feast.infra.offline_stores.contrib.spark_offline_store.uc_registration.get_databricks_session" - ) as mock_get_session: - mock_session = MagicMock() - mock_get_session.return_value = mock_session - - config = DatabricksUCOfflineStoreConfig( - type="databricks_uc", - default_catalog=None, - default_schema=None, - uc_registration=UCRegistrationConfig(enabled=True), - ) - - fv = MagicMock(spec=FeatureView) - fv.tags = {} - fv.name = "no_cat_fv" - fv.entity_columns = [] - fv.features = [] - fv.batch_source = MagicMock() - fv.batch_source.timestamp_field = None - fv.description = "" - fv.owner = None - - register_uc_feature_tables(config, [fv], "proj") + "feast.infra.offline_stores.iceberg.registration.update_table_properties" + ) as mock_update: + register_uc_feature_tables(config_no_warehouse, [fv], "proj") - fe_client.create_table.assert_not_called() + mock_create.assert_not_called() + mock_update.assert_not_called() From 4f41ccf39fa94ab22c6c3bff76f015239fd380fa Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sun, 28 Jun 2026 15:09:29 +0530 Subject: [PATCH 07/16] feat: Implement UC-backed materialization via compute engine hooks (L3) Signed-off-by: Abhishek Shinde --- .../feast/infra/compute_engines/local/nodes.py | 12 ++++++++++++ .../feast/infra/compute_engines/spark/nodes.py | 12 ++++++++++++ .../spark_offline_store/test_uc_registration.py | 2 ++ 3 files changed, 26 insertions(+) diff --git a/sdk/python/feast/infra/compute_engines/local/nodes.py b/sdk/python/feast/infra/compute_engines/local/nodes.py index 9d3e1a48881..2e37863db4b 100644 --- a/sdk/python/feast/infra/compute_engines/local/nodes.py +++ b/sdk/python/feast/infra/compute_engines/local/nodes.py @@ -407,4 +407,16 @@ def execute(self, context: ExecutionContext) -> ArrowTableValue: progress=lambda x: None, ) + # UC-backed materialization hook (Phase L3) + from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( + write_uc_materialized_data, + ) + + write_uc_materialized_data( + config=context.repo_config, + fv=self.feature_view, + df=input_table, + project=context.repo_config.project, + ) + return output diff --git a/sdk/python/feast/infra/compute_engines/spark/nodes.py b/sdk/python/feast/infra/compute_engines/spark/nodes.py index 92964b72bc9..0df3b774973 100644 --- a/sdk/python/feast/infra/compute_engines/spark/nodes.py +++ b/sdk/python/feast/infra/compute_engines/spark/nodes.py @@ -590,6 +590,18 @@ def execute(self, context: ExecutionContext) -> DAGValue: ) spark_df.write.format(file_format).mode("append").save(dest_path) + # UC-backed materialization hook (Phase L3) + from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( + write_uc_materialized_data, + ) + + write_uc_materialized_data( + config=context.repo_config, + fv=self.feature_view, + df=spark_df, + project=context.repo_config.project, + ) + return DAGValue( data=spark_df, format=DAGFormat.SPARK, diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py index cd0ebe6adee..0c72a3eac34 100644 --- a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py @@ -9,7 +9,9 @@ _resolve_uc_path, _should_register, register_uc_feature_tables, + write_uc_materialized_data, ) +from feast.repo_config import RepoConfig def make_mock_field(name: str, dtype_str: str = "", nullable: bool = True): From c2a0d271593883aa04b436e4dea59164c51e27c8 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Wed, 1 Jul 2026 14:32:35 -0700 Subject: [PATCH 08/16] feat: Add click-to-zoom lightbox for blog post images (#6575) Clicking any image in a blog post now opens a fullscreen overlay with the image centered on a dark backdrop. Close with click or Escape. Signed-off-by: Francisco Javier Arceo Co-authored-by: Claude Opus 4.6 --- infra/website/src/pages/blog/[slug].astro | 66 ++++++++++++++++++++++- 1 file changed, 65 insertions(+), 1 deletion(-) diff --git a/infra/website/src/pages/blog/[slug].astro b/infra/website/src/pages/blog/[slug].astro index 90e8572d0b2..a8ab3da8559 100644 --- a/infra/website/src/pages/blog/[slug].astro +++ b/infra/website/src/pages/blog/[slug].astro @@ -223,6 +223,7 @@ const socialImage = heroImage ? toAbsoluteUrl(heroImage) : undefined; height: auto; border-radius: 4px; margin: 32px 0; + cursor: zoom-in; } /* Hero Image */ @@ -339,4 +340,67 @@ const socialImage = heroImage ? toAbsoluteUrl(heroImage) : undefined; font-size: 18px; } } - + + + + + + From 9446524f3379124955a4a8c72d0710b5cc0625b3 Mon Sep 17 00:00:00 2001 From: Aniket Paluskar Date: Wed, 1 Jul 2026 13:58:08 +0530 Subject: [PATCH 09/16] fix: configure FIPS-compliant gRPC cipher suites for offline server Signed-off-by: Aniket Paluska Co-authored-by: Cursor Signed-off-by: Aniket Paluskar --- sdk/python/feast/offline_server.py | 26 +++++++++ sdk/python/tests/unit/test_offline_server.py | 55 +++++++++++++++++++- 2 files changed, 79 insertions(+), 2 deletions(-) diff --git a/sdk/python/feast/offline_server.py b/sdk/python/feast/offline_server.py index 0373b3a8146..cdf3d9edaa6 100644 --- a/sdk/python/feast/offline_server.py +++ b/sdk/python/feast/offline_server.py @@ -38,6 +38,31 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) +_FIPS_CIPHER_SUITES = ":".join( + [ + "ECDHE-RSA-AES128-GCM-SHA256", + "ECDHE-RSA-AES256-GCM-SHA384", + "ECDHE-ECDSA-AES128-GCM-SHA256", + "ECDHE-ECDSA-AES256-GCM-SHA384", + "AES128-GCM-SHA256", + "AES256-GCM-SHA384", + ] +) + + +def _is_fips_enabled() -> bool: + try: + with open("/proc/sys/crypto/fips_enabled") as f: + return f.read().strip() == "1" + except (FileNotFoundError, PermissionError, OSError): + return False + + +def _configure_grpc_fips() -> None: + if _is_fips_enabled() and "GRPC_SSL_CIPHER_SUITES" not in os.environ: + os.environ["GRPC_SSL_CIPHER_SUITES"] = _FIPS_CIPHER_SUITES + logger.info("FIPS mode detected, configured FIPS-compliant gRPC cipher suites.") + class OfflineServer(fl.FlightServerBase): def __init__( @@ -596,6 +621,7 @@ def start_server( tls_cert_path: str = "", ): _init_auth_manager(store) + _configure_grpc_fips() tls_certificates = [] scheme = "grpc+tcp" diff --git a/sdk/python/tests/unit/test_offline_server.py b/sdk/python/tests/unit/test_offline_server.py index 3e25e5c2061..6ff93d02282 100644 --- a/sdk/python/tests/unit/test_offline_server.py +++ b/sdk/python/tests/unit/test_offline_server.py @@ -1,4 +1,5 @@ -from unittest.mock import MagicMock, patch +import os +from unittest.mock import MagicMock, mock_open, patch import assertpy @@ -7,7 +8,11 @@ RemoteOfflineStoreConfig, _create_retrieval_metadata, ) -from feast.offline_server import OfflineServer +from feast.offline_server import ( + OfflineServer, + _configure_grpc_fips, + _is_fips_enabled, +) def test_create_retrieval_metadata_with_sql_string(): @@ -86,3 +91,49 @@ def test_offline_server_get_historical_features_passes_sql_to_store(): assertpy.assert_that(result).is_equal_to(mock_job) _, kwargs = mock_offline_store.get_historical_features.call_args assertpy.assert_that(kwargs["entity_df"]).is_equal_to(sql) + + +def test_is_fips_enabled_returns_true(): + with patch("builtins.open", mock_open(read_data="1\n")): + assert _is_fips_enabled() is True + + +def test_is_fips_enabled_returns_false(): + with patch("builtins.open", mock_open(read_data="0\n")): + assert _is_fips_enabled() is False + + +def test_is_fips_enabled_missing_file(): + with patch("builtins.open", side_effect=FileNotFoundError): + assert _is_fips_enabled() is False + + +def test_configure_grpc_fips_sets_cipher_suites(): + with ( + patch("feast.offline_server._is_fips_enabled", return_value=True), + patch.dict(os.environ, {}, clear=False), + ): + os.environ.pop("GRPC_SSL_CIPHER_SUITES", None) + _configure_grpc_fips() + assert "GRPC_SSL_CIPHER_SUITES" in os.environ + assert "AES128-GCM-SHA256" in os.environ["GRPC_SSL_CIPHER_SUITES"] + del os.environ["GRPC_SSL_CIPHER_SUITES"] + + +def test_configure_grpc_fips_respects_existing_env(): + with ( + patch("feast.offline_server._is_fips_enabled", return_value=True), + patch.dict(os.environ, {"GRPC_SSL_CIPHER_SUITES": "custom"}, clear=False), + ): + _configure_grpc_fips() + assert os.environ["GRPC_SSL_CIPHER_SUITES"] == "custom" + + +def test_configure_grpc_fips_noop_without_fips(): + with ( + patch("feast.offline_server._is_fips_enabled", return_value=False), + patch.dict(os.environ, {}, clear=False), + ): + os.environ.pop("GRPC_SSL_CIPHER_SUITES", None) + _configure_grpc_fips() + assert "GRPC_SSL_CIPHER_SUITES" not in os.environ From 5ce0cf72b0a442cb409134d306c87ff7fbe22e89 Mon Sep 17 00:00:00 2001 From: Valentyn Kahamlyk Date: Wed, 1 Jul 2026 23:04:51 -0700 Subject: [PATCH 10/16] feat: Add OnlineStore for Aerospike (#6532) * feat: scaffold Aerospike online store Signed-off-by: Valentyn Kahamlyk * feat: implement Aerospike online_write_batch Signed-off-by: Valentyn Kahamlyk * feat: implement Aerospike online_read Signed-off-by: Valentyn Kahamlyk * feat: implement Aerospike update and teardown Signed-off-by: Valentyn Kahamlyk * test: add Aerospike unit and integration tests Signed-off-by: Valentyn Kahamlyk * feat: add async online_read/write and lifecycle hooks for Aerospike Signed-off-by: Valentyn Kahamlyk * docs: add Aerospike online store reference and tuning guide Signed-off-by: Valentyn Kahamlyk * fix: use bytearray keys and zip-based batch mapping for Aerospike reads The Aerospike Python client mishandles bytes user keys (hashes only the first byte), collapsing all entities onto the same digest. Wrap keys in bytearray on write and read. Also pair BatchRecord responses with original input keys via zip rather than trusting br.key[2], which the client returns in a different representation on reads. Add two integration tests: cross-FV Map CDT coexistence and update(tables_to_delete=...) background scan. Signed-off-by: Valentyn Kahamlyk * docs: clarify Aerospike auth and TLS sections are Enterprise-only Signed-off-by: Valentyn Kahamlyk * feat: add aerospike to feast-operator supported online stores Signed-off-by: Valentyn Kahamlyk * fix(aerospike): project requested_features server-side and surface per-record errors Two review blockers rolled into one commit because they share the same code path. 1. Server-side projection. online_read now builds a map_get_by_key_list op nested into the feature-view submap via cdt_ctx_map_key when requested_features is provided, instead of fetching the whole FV slot and filtering in Python. For wide feature views this ships only the requested columns over the wire. The response shape (flat [k,v,k,v] list vs. dict) is normalized through _normalize_projected_features. 2. Per-record error surfacing. Both batch_write and batch_operate only raise when the whole request is rejected; partial failures (single-partition timeout, replica quorum miss) are otherwise silent and present downstream as missing features. online_read now distinguishes RECORD_NOT_FOUND (2) and OP_NOT_APPLICABLE (26, = nested ctx miss when FV slot is absent) from transient errors, which are raised. online_write_batch inspects every per-record result code after the batch call. Unit tests cover all four paths: projected read, not-found, op-not-applicable (nested ctx miss), and a simulated TIMEOUT that must raise. The docker-backed cross-FV and update() integration tests still pass, so server-side projection is verified end-to-end against a real Aerospike server. Signed-off-by: Valentyn Kahamlyk * feat(aerospike)!: rename total_timeout_ms -> batch_total_timeout_ms and add socket_timeout_ms Review feedback: total_timeout_ms was ambiguous (users read it as a global/end-to-end timeout) and the timeout surface was missing socket_timeout, which is the per-attempt trigger that lets max_retries actually fire within the total budget. * total_timeout_ms -> batch_total_timeout_ms. Now explicitly named after the Aerospike batch policy it maps to, matches read_timeout_ms / write_timeout_ms in framing (each targets one policy scope). * Add socket_timeout_ms (optional). Applies uniformly to read, write and batch policies when set. Leaves the Aerospike client default in place when unset. BREAKING CHANGE: total_timeout_ms is renamed to batch_total_timeout_ms. Config files using the old name must be updated. No default value change. Docs updated (reference + perf-tuning guide) with a short explainer on the per-attempt vs total deadline distinction. Two new unit tests pin the policy wiring: socket_timeout_ms propagates to all three scopes, and is omitted (not injected as None) when unset. Signed-off-by: Valentyn Kahamlyk * refactor(aerospike): use MAP_KEY_ORDERED, KEY_DIGEST, and instance-scoped client Cheap-win cleanups flagged in review, all touching the same small patch of write-path and lifecycle code. * Map CDTs are now created with MAP_KEY_ORDERED. map_get_by_key / map_remove_by_key on an ordered map are O(log N) in the map size instead of O(N); matters on reads of wide feature views and on the update() background scan (which walks every record in the project's set). * Writes drop POLICY_KEY_SEND and rely on the client default (POLICY_KEY_DIGEST). The serialized entity key is no longer stored alongside each record, saving per-record storage the read path never consumes (batch_operate preserves request order; results are paired back by zip in online_read). * _client moves from a class attribute to an instance attribute (set in __init__). Previously two AerospikeOnlineStore instances could share the cached client through class state until one wrote self._client. With the instance attribute the state is always per-instance from construction. * Drop MongoDB references from class docstrings and comments (they referred to how the storage layout was derived rather than documenting current behavior). Also rewrite the _build_batch_writes docstring to describe the policies applied on the write path. Unit test assertions for the write-path record are updated: bw.policy is now None (client default applies) and map ops carry map_policy={'map_order': MAP_KEY_ORDERED}. All three docker-backed integration tests still pass end-to-end (cross-FV upsert, update() background scan, full feature-store round-trip), so the read/write shape survives the ordering and policy changes against a real server. Signed-off-by: Valentyn Kahamlyk * feat(aerospike): add per-FV namespace/set overrides and prewriting hook Adds three configuration knobs to AerospikeOnlineStoreConfig: - namespace_overrides: pin individual feature views to a different Aerospike namespace (e.g. RAM-only vs. SSD-backed) without splitting the project across stores. - set_overrides: place a feature view in its own set so admin ops on it (truncate, scan-based deletes during `feast apply`) do not touch records of other views. - prewriting_hook: import-string-resolved callable invoked once per online_write_batch with the rows about to be written, returning the rows that actually go on the wire. Resolved and cached on first use; returning [] short-circuits the wire call. Read, write, update and teardown paths all honour the per-FV ns/set resolution. update() groups dropped feature views by their resolved (ns, set) pair and issues one background scan per group. teardown() truncates every unique (ns, set) pair the project may have written to, including the store-level default. Adds 22 unit tests for the new behaviour and updates 3 existing call sites of _build_batch_writes for the new namespace= parameter. Adds a sample hook module under examples/online_store/aerospike_overrides_and_hooks/ and corresponding sections in docs/reference/online-stores/aerospike.md. Signed-off-by: Valentyn Kahamlyk * test: update aerospike image tag Signed-off-by: Valentyn Kahamlyk * chore: sync README template and secrets baseline after master merge Signed-off-by: Valentyn Kahamlyk * chore: fix secrets baseline line number for v1 operator types Adding aerospike to the feast-operator enum shifted the allowlisted SecretRef entry in api/v1/featurestore_types.go by one line. Signed-off-by: Valentyn Kahamlyk * docs: update aerospike docs Signed-off-by: Valentyn Kahamlyk * fix(aerospike): wire batch max_retries and fix empty projection handling Copilot review feedback on PR #6532: - Add max_retries to the batch client policy (batch_operate/batch_write path) - Treat empty projected feature maps as present FV slots (is not None) - Return {} from _normalize_projected_features([]) instead of None - Fix projection unit test mock/assertions - Correct prewriting_hook config docstring Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Valentyn Kahamlyk * style(aerospike): format online_read docs assignment for ruff Signed-off-by: Valentyn Kahamlyk * chore: update pixi.lock for aerospike optional extra Regenerate the v6 lockfile with Pixi v0.63.1 after adding the aerospike extra to pyproject.toml. Signed-off-by: Valentyn Kahamlyk * fix(aerospike): add client init lock and batch chunking Guard lazy client creation with a lock to avoid connection leaks under concurrent first use, and chunk batch reads/writes by batch_max_records so large materializations stay under Aerospike server batch limits. Signed-off-by: Valentyn Kahamlyk --------- Signed-off-by: Valentyn Kahamlyk Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .secrets.baseline | 4 +- README.md | 1 + docs/SUMMARY.md | 1 + .../online-server-performance-tuning.md | 31 + docs/reference/online-stores/README.md | 4 + docs/reference/online-stores/aerospike.md | 389 ++++ docs/roadmap.md | 1 + .../aerospike_overrides_and_hooks/README.md | 65 + .../feature_store.yaml | 59 + .../aerospike_overrides_and_hooks/hooks.py | 109 ++ .../api/v1/featurestore_types.go | 3 +- .../api/v1alpha1/featurestore_types.go | 3 +- .../manifests/feast.dev_featurestores.yaml | 4 + .../crd/bases/feast.dev_featurestores.yaml | 4 + infra/feast-operator/dist/install.yaml | 4 + pixi.lock | 5 +- pyproject.toml | 3 + .../aerospike_online_store/__init__.py | 3 + .../aerospike_online_store/aerospike.py | 1047 ++++++++++ .../aerospike_repo_configuration.py | 13 + sdk/python/feast/repo_config.py | 1 + .../test_aerospike_online_retrieval.py | 1684 +++++++++++++++++ .../universal/online_store/aerospike.py | 58 + 23 files changed, 3490 insertions(+), 6 deletions(-) create mode 100644 docs/reference/online-stores/aerospike.md create mode 100644 examples/online_store/aerospike_overrides_and_hooks/README.md create mode 100644 examples/online_store/aerospike_overrides_and_hooks/feature_store.yaml create mode 100644 examples/online_store/aerospike_overrides_and_hooks/hooks.py create mode 100644 sdk/python/feast/infra/online_stores/aerospike_online_store/__init__.py create mode 100644 sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike.py create mode 100644 sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike_repo_configuration.py create mode 100644 sdk/python/tests/unit/online_store/test_aerospike_online_retrieval.py create mode 100644 sdk/python/tests/universal/feature_repos/universal/online_store/aerospike.py diff --git a/.secrets.baseline b/.secrets.baseline index 1ba1359f336..0eefabade70 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -957,7 +957,7 @@ "filename": "infra/feast-operator/api/v1/featurestore_types.go", "hashed_secret": "44e17306b837162269a410204daaa5ecee4ec22c", "is_verified": false, - "line_number": 935 + "line_number": 936 } ], "infra/feast-operator/api/v1/zz_generated.deepcopy.go": [ @@ -989,7 +989,7 @@ "filename": "infra/feast-operator/api/v1alpha1/featurestore_types.go", "hashed_secret": "44e17306b837162269a410204daaa5ecee4ec22c", "is_verified": false, - "line_number": 649 + "line_number": 650 } ], "infra/feast-operator/api/v1alpha1/zz_generated.deepcopy.go": [ diff --git a/README.md b/README.md index 0510ba42f1f..3b35344b021 100644 --- a/README.md +++ b/README.md @@ -227,6 +227,7 @@ The list below contains the functionality that contributors are planning to deve * [x] [SingleStore](https://docs.feast.dev/reference/online-stores/singlestore) * [x] [Couchbase](https://docs.feast.dev/reference/online-stores/couchbase) * [x] [MongoDB](https://docs.feast.dev/reference/online-stores/mongodb) + * [x] [Aerospike](https://docs.feast.dev/reference/online-stores/aerospike) * [x] [Qdrant (vector store)](https://docs.feast.dev/reference/online-stores/qdrant) * [x] [Milvus (vector store)](https://docs.feast.dev/reference/online-stores/milvus) * [x] [Faiss (vector store)](https://docs.feast.dev/reference/online-stores/faiss) diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md index 4e39561c10f..547c88acf68 100644 --- a/docs/SUMMARY.md +++ b/docs/SUMMARY.md @@ -164,6 +164,7 @@ * [SingleStore](reference/online-stores/singlestore.md) * [Milvus](reference/online-stores/milvus.md) * [MongoDB](reference/online-stores/mongodb.md) + * [Aerospike](reference/online-stores/aerospike.md) * [Elasticsearch](reference/online-stores/elasticsearch.md) * [Qdrant](reference/online-stores/qdrant.md) * [Faiss](reference/online-stores/faiss.md) diff --git a/docs/how-to-guides/online-server-performance-tuning.md b/docs/how-to-guides/online-server-performance-tuning.md index b280fe53cff..f10d7bff8c8 100644 --- a/docs/how-to-guides/online-server-performance-tuning.md +++ b/docs/how-to-guides/online-server-performance-tuning.md @@ -278,6 +278,7 @@ The online store is the single largest factor in `get_online_features()` latency | **DynamoDB** | 2–5 ms | Yes | Serverless, auto-scaling on AWS | Pay-per-request cost; batch API limits (100 items) | | **PostgreSQL** | 3–10 ms | No (threadpool) | Teams with existing Postgres infra | Connection pooling needed at scale | | **MongoDB** | 2–5 ms | Yes | Flexible schema, async-native | Requires index tuning for large datasets | +| **Aerospike** | < 1 ms | No (threadpool) | Ultra-low latency, hybrid memory (RAM + SSD), large datasets | Namespace must be pre-configured on the cluster | | **Bigtable** | 3–8 ms | No (threadpool) | Large-scale GCP workloads | Row-key design affects read performance | | **Cassandra / ScyllaDB** | 2–5 ms | No (threadpool) | Multi-region, write-heavy | Tunable consistency; requires DC-aware routing | | **Remote** | Varies | No (threadpool) | Centralized feature server architecture | Adds an HTTP hop; tune connection pool | @@ -305,6 +306,7 @@ The feature server can read from the online store using either an **async** or * | **MongoDB** | Yes | Yes | Uses `motor` (async MongoDB driver) | | **PostgreSQL** | Implemented | No | Has `online_read_async` but does not yet advertise via `async_supported`; uses sync/threadpool path | | **Redis** | Implemented | **Yes** | `online_read_async` and `online_write_batch_async` both implemented; uses sync/threadpool path for `get_online_features` (overridden with batched single pipeline) | +| **Aerospike** | Implemented | No | Async methods wrap the blocking C client via `run_in_executor`; does not yet advertise via `async_supported`, so the server still uses the threadpool path | | All others | No | No | Fall back to sync with `run_in_threadpool()` | **When async matters most:** @@ -467,6 +469,34 @@ online_store: - **`connectTimeoutMS` / `socketTimeoutMS`**: Tighter timeouts improve p99 by failing fast on slow connections. - MongoDB is one of the stores with **full async support** (read and write), so it benefits from concurrent feature view reads via `asyncio.gather()`. +### Aerospike tuning + +Aerospike offers sub-millisecond reads thanks to its hybrid-memory architecture (primary index in RAM, data on SSD or RAM). Tune the per-call policies in the Feast config and rely on the Aerospike cluster's own tuning for everything else: + +```yaml +online_store: + type: aerospike + hosts: + - ["aerospike-1.internal", 3000] + - ["aerospike-2.internal", 3000] + namespace: feast + read_timeout_ms: 150 # hard deadline for a single-record get + write_timeout_ms: 300 # hard deadline for a single-record put/operate + batch_total_timeout_ms: 500 # hard deadline for online_read / online_write_batch + socket_timeout_ms: 50 # per-attempt deadline so max_retries can actually fire + max_retries: 2 + ttl_seconds: 86400 # record-level TTL; omit to use the namespace default + client_kwargs: # escape hatch for any client-config field not surfaced above + policies: + batch: + concurrent_nodes: 0 # 0 = parallel to every node (lowest latency on multi-node clusters) +``` + +- **`*_timeout_ms` (total)** vs **`socket_timeout_ms` (per-attempt)**: `*_timeout_ms` is the hard deadline for a whole call *including* retries; `socket_timeout_ms` is the per-attempt deadline that allows `max_retries` to actually fire within that budget. Without `socket_timeout_ms`, a single slow attempt can consume the entire total deadline and retries never run. +- **`hosts`**: List every seed node. The Aerospike client discovers the rest of the cluster automatically and opens one connection pool per node. +- **`ttl_seconds: 0`** means "never expire"; omit the key to inherit the namespace's `default-ttl`. Expiry is enforced by the server's `nsup` thread — nothing to delete on the client side. +- Co-locate the feature server in the **same availability zone / rack** as the Aerospike cluster; sub-millisecond reads are bandwidth- and RTT-sensitive. + ### Remote online store tuning The Remote online store connects to a Feast feature server over HTTP. Connection pooling is critical: @@ -700,6 +730,7 @@ This applies to every connection-oriented online store: | **DynamoDB** | `max_pool_connections` (HTTP pool) | 10 | No hard limit, but AWS SDK has per-process pool caps; monitor throttling | | **Redis** | Connection per worker | 1 | `maxclients` on the Redis server (default: 10,000) | | **MongoDB** | `maxPoolSize` (in `client_kwargs`) | 100 | Server's `net.maxIncomingConnections` | +| **Aerospike** | Driver manages pool per seed node | Auto | `proto-fd-max` (default 15000) on each Aerospike node | | **Cassandra** | Driver manages pool per node | Auto | `native_transport_max_threads` on each Cassandra node | | **Remote** | `connection_pool_size` (HTTP pool) | 50 | The target feature server's worker capacity | diff --git a/docs/reference/online-stores/README.md b/docs/reference/online-stores/README.md index 6f31993f896..39294966170 100644 --- a/docs/reference/online-stores/README.md +++ b/docs/reference/online-stores/README.md @@ -58,6 +58,10 @@ Please see [Online Store](../../getting-started/components/online-store.md) for [mongodb.md](mongodb.md) {% endcontent-ref %} +{% content-ref url="aerospike.md" %} +[aerospike.md](aerospike.md) +{% endcontent-ref %} + {% content-ref url="hazelcast.md" %} [hazelcast.md](hazelcast.md) {% endcontent-ref %} diff --git a/docs/reference/online-stores/aerospike.md b/docs/reference/online-stores/aerospike.md new file mode 100644 index 00000000000..e5a9754796b --- /dev/null +++ b/docs/reference/online-stores/aerospike.md @@ -0,0 +1,389 @@ +# Aerospike online store (Preview) + +## Description + +The [Aerospike](https://aerospike.com/) online store provides support for materializing feature values into an Aerospike cluster for serving online features. + +{% hint style="warning" %} +The Aerospike online store is currently in **preview**. Some functionality may be unstable, and breaking changes may occur in future releases. +{% endhint %} + +## Features + +* Supports both synchronous and asynchronous read/write paths (`online_read` / `online_read_async`, `online_write_batch` / `online_write_batch_async`). Async methods wrap the blocking client in `run_in_executor`, keeping the event loop responsive in feature-server workloads. +* Partial, server-side upserts via Aerospike Map CDT operations — writing one feature view never clobbers another feature view stored on the same entity. +* Record-level TTL controlled by a single `ttl_seconds` config option (honours the namespace default, a "never expire" sentinel, or an explicit number of seconds). +* Per-feature-view **namespace overrides** and **set overrides** — pin individual feature views to RAM-only or SSD-backed namespaces, or isolate one view in its own set, without splitting projects. +* **Prewriting hook** — a configurable, import-string-resolved callable applied to every write batch for cross-cutting concerns like PII masking, application-side encryption, or value coercion. +* Authentication and TLS options for Aerospike Enterprise Edition passed straight through to the Aerospike Python client. +* `client_kwargs` escape hatch for any advanced client-config field not surfaced on `AerospikeOnlineStoreConfig`. +* Baseline: Aerospike Server **≥ 6.0** (uses batch-write / batch-operate APIs). The store has been developed against CE 8.x. + +## Getting started + +Install the Aerospike extra (alongside the dependency for the offline store of choice): + +```bash +pip install 'feast[aerospike]' +``` + +You can start from any of the standard templates (e.g. `feast init -t local` or `feast init -t aws`) and then swap in Aerospike as the online store as shown below. + +## Examples + +### Basic configuration — local Aerospike CE + +{% code title="feature_store.yaml" %} +```yaml +project: my_feature_repo +registry: data/registry.db +provider: local +online_store: + type: aerospike + hosts: + - ["127.0.0.1", 3000] + namespace: feast +``` +{% endcode %} + +### Multi-node cluster + +{% code title="feature_store.yaml" %} +```yaml +project: my_feature_repo +registry: data/registry.db +provider: local +online_store: + type: aerospike + hosts: + - ["aerospike-1.internal", 3000] + - ["aerospike-2.internal", 3000] + - ["aerospike-3.internal", 3000] + namespace: feast + ttl_seconds: 86400 # 24h record-level TTL + read_timeout_ms: 150 # hard deadline for a single-record get + write_timeout_ms: 300 # hard deadline for a single-record put/operate + batch_total_timeout_ms: 500 # hard deadline for online_read / online_write_batch + batch_max_records: 1000 # chunk size for batch_write / batch_operate + socket_timeout_ms: 50 # per-attempt deadline so max_retries can fire + max_retries: 2 +``` +{% endcode %} + +> **Timeout semantics.** The Aerospike client distinguishes per-attempt +> (`socket_timeout`) from total (`total_timeout`) deadlines. `*_timeout_ms` map +> to `total_timeout` — the overall budget for a call including retries. Set +> `socket_timeout_ms` as well so each individual attempt has its own (shorter) +> deadline; without it, `max_retries` effectively never fires because the +> first attempt is allowed to consume the entire total deadline. + +> **Batch chunking.** `online_read` and `online_write_batch` split large +> requests into chunks of at most `batch_max_records` (default `1000`). +> Aerospike enforces a per-node batch limit via the server `batch-max-requests` +> setting (historically `5000`). Lower `batch_max_records` if your cluster cap +> is tighter; raise it only when the server limit and client timeouts allow. + +### Aerospike Enterprise with authentication + +> Requires Aerospike Enterprise Edition. The Community Edition server has no built-in user/security model and will reject these config keys. + +{% code title="feature_store.yaml" %} +```yaml +project: my_feature_repo +registry: data/registry.db +provider: local +online_store: + type: aerospike + hosts: + - ["aerospike.internal", 3000] + namespace: feast + user: feast_user + password: ${AEROSPIKE_PASSWORD} # pragma: allowlist secret + auth_mode: internal # internal | external | pki +``` +{% endcode %} + +### Aerospike Enterprise with TLS + +> Requires Aerospike Enterprise Edition. The Community Edition server does not implement TLS, so `tls` config is effective only against EE clusters. + +{% code title="feature_store.yaml" %} +```yaml +project: my_feature_repo +registry: data/registry.db +provider: local +online_store: + type: aerospike + hosts: + - ["aerospike-1.internal", 4333, "aerospike-tls"] + namespace: feast + tls: + enable: true + cafile: /etc/aerospike/certs/ca.pem + certfile: /etc/aerospike/certs/client.pem + keyfile: /etc/aerospike/certs/client.key +``` +{% endcode %} + +### Per-feature-view namespace and set overrides + +Two `Dict[str, str]` config fields — `namespace_overrides` and `set_overrides` — let you place individual feature views on a different Aerospike namespace or set without splitting your project across stores. Anything not listed in either map falls back to the store-level default (`namespace` / `set_name_template`). + +Common reasons to reach for these: + +* A **hot, latency-sensitive view** belongs on a RAM-only namespace; a **wide, cold view** belongs on an SSD-backed namespace. Same project, different storage tiers. +* You want `feast apply` deletions or `truncate` on one feature view to be O(1) without scanning records of the others — give that view its own set. + +{% code title="feature_store.yaml" %} +```yaml +project: my_feature_repo +registry: data/registry.db +provider: local +online_store: + type: aerospike + hosts: + - ["aerospike.internal", 3000] + namespace: feast # default namespace + set_name_template: "{project}_{collection_suffix}" + namespace_overrides: + driver_realtime_stats: feast_ram # in-memory namespace + driver_history_lookup: feast_ssd # device-backed namespace + set_overrides: + isolated_view: my_feature_repo_isolated +``` +{% endcode %} + +> **Tradeoffs.** +> +> * Every namespace listed in `namespace_overrides` MUST already exist on the cluster — Aerospike cannot create namespaces at runtime, and a missing namespace surfaces as an opaque `AEROSPIKE_ERR_PARAM` on the first read or write. +> * Putting feature views on different sets means a multi-feature-view read for the same entity becomes one Aerospike round trip per set, not one round trip total. Only opt in when the operational isolation is worth that cost. Reads that touch a single feature view are unaffected. +> * Admin operations honour the overrides automatically: `update()` (called by `feast apply`) groups dropped feature views by their resolved `(namespace, set)` and issues one background scan per group; `teardown()` truncates every unique `(namespace, set)` pair the project may have written to (including the store-level default). + +### Prewriting hooks + +`prewriting_hook` is the import path of a callable that is invoked once per `online_write_batch` call, receives the rows about to be written, and returns the rows that actually go on the wire. Use it for cross-cutting write-side concerns that you don't want sprinkled through every materialization job — PII masking, application-side encryption, dual-write fan-out, value coercion, etc. + +Hooks are referenced by import string (rather than as a Python `Callable` value) so the config survives YAML/JSON serialisation and remote-feature-server transport. The resolved callable is cached on the store instance, so import cost is paid once per store lifetime. + +**Hook signature:** + +```python +def hook( + config: RepoConfig, + table: FeatureView, + data: list[ + tuple[ + EntityKeyProto, + dict[str, ValueProto], + datetime, + datetime | None, + ] + ], +) -> list[ + tuple[ + EntityKeyProto, + dict[str, ValueProto], + datetime, + datetime | None, + ] +]: + ... +``` + +The hook MUST return a row list with the same schema as its input. Returning `[]` short-circuits the write — same path as an empty input, no wire call is issued. Hooks that raise will fail the whole batch; there is no per-row fallback. + +**1. Drop a hook function in your project.** Any module on the `PYTHONPATH` of every process that writes through Feast will do (the materialization workers, the registry CLI host, and the feature server, if you run one). + +{% code title="my_feature_repo/hooks.py" %} +```python +"""Prewriting hooks for the Aerospike online store.""" +from __future__ import annotations + +import hashlib +import os +from datetime import datetime +from typing import Optional + +from feast import FeatureView +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto +from feast.repo_config import RepoConfig + +# Names of features that must never reach the online store as plaintext. +# Matched by exact feature name; tweak to your project's conventions. +_SENSITIVE_FEATURES = {"email", "phone_number", "ssn"} + + +def hash_pii_string_features( + config: RepoConfig, + table: FeatureView, + data: list[ + tuple[ + EntityKeyProto, + dict[str, ValueProto], + datetime, + Optional[datetime], + ] + ], +) -> list[ + tuple[ + EntityKeyProto, + dict[str, ValueProto], + datetime, + Optional[datetime], + ] +]: + """Replace any sensitive string feature with a salted SHA-256 hex digest. + + The hash is deterministic (same input → same digest) so downstream lookups + that hash the candidate value the same way still hit. ``FEAST_PII_SALT`` + must be set on every process that materialises features; an unset salt + raises rather than silently falling back to plaintext. + """ + salt = os.environ.get("FEAST_PII_SALT") + if salt is None: + raise RuntimeError( + "FEAST_PII_SALT is not set; refusing to write feature batches " + "without a configured PII salt." + ) + salt_bytes = salt.encode("utf-8") + + def _digest(plaintext: str) -> str: + h = hashlib.sha256() + h.update(salt_bytes) + h.update(plaintext.encode("utf-8")) + return h.hexdigest() + + transformed: list[ + tuple[ + EntityKeyProto, + dict[str, ValueProto], + datetime, + Optional[datetime], + ] + ] = [] + for entity_key, values, event_ts, created_ts in data: + new_values = dict(values) + for feature_name in _SENSITIVE_FEATURES.intersection(new_values): + v = new_values[feature_name] + if v.HasField("string_val") and v.string_val: + new_values[feature_name] = ValueProto(string_val=_digest(v.string_val)) + transformed.append((entity_key, new_values, event_ts, created_ts)) + return transformed +``` +{% endcode %} + +**2. Reference the hook from `feature_store.yaml`:** + +{% code title="feature_store.yaml" %} +```yaml +project: my_feature_repo +registry: data/registry.db +provider: local +online_store: + type: aerospike + hosts: + - ["aerospike.internal", 3000] + namespace: feast + prewriting_hook: my_feature_repo.hooks.hash_pii_string_features +``` +{% endcode %} + +> **Operational notes.** +> +> * The hook is **only invoked on the write path**; reads pass through the store untouched. If your hook is one-way (e.g. hashing) you have to apply the same transformation to the candidate value at read time yourself. +> * Hooks run inside the same process as the writer — they're not RPCs and not sandboxed. They can read environment variables, open files, call out to KMS, etc. Treat them as part of your trusted code base. +> * A misconfigured `prewriting_hook` (bad import path, missing function, non-callable target) raises `ValueError` / `TypeError` on the *first* `online_write_batch` call, not on store construction. Add a smoke test that writes one row at deploy time so misconfigurations surface before a real batch. + +The full set of configuration options is available in [`AerospikeOnlineStoreConfig`](https://rtd.feast.dev/en/latest/#feast.infra.online_stores.aerospike_online_store.aerospike.AerospikeOnlineStoreConfig). + +## Data Model + +The Aerospike online store uses a **single set per project** with entity-key collocation. Features from multiple feature views for the same entity are stored together on a single Aerospike record, analogous to the MongoDB online store's "one document per entity" layout. + +| Aerospike concept | Feast mapping | +| :---------------- | :---------------------------------------------------------------------------- | +| Namespace | `online_store.namespace` (must be pre-configured on the cluster); per-feature-view override via `online_store.namespace_overrides` | +| Set | `online_store.set_name_template` → `"{project}_{collection_suffix}"` by default; per-feature-view override via `online_store.set_overrides` | +| Key | `serialize_entity_key(entity_key)` as `bytearray` user key | +| Bin `features` | Map CDT keyed by feature-view name, each value a map of `feature → native` | +| Bin `event_ts` | Map CDT keyed by feature-view name, each value an int64 epoch-ms timestamp | +| Bin `created_ts` | Top-level int64 epoch-ms timestamp (last `feast materialize`) | + +### Example record + +For a single entity carrying features from two feature views (`driver_stats` and `pricing`): + +```text +key: (ns="feast", set="my_feature_repo_latest", user_key=) +bins: + features: + driver_stats: + rating: 4.91 + trips_last_7d: 132 + pricing: + surge_multiplier: 1.2 + event_ts: + driver_stats: 1737374400000 # 2025-01-20T12:00:00Z + pricing: 1737447000000 # 2025-01-21T08:30:00Z + created_ts: 1737460805000 # 2025-01-21T12:00:05Z +``` + +### Key design decisions + +* **Record per entity, bin per concept.** `features` and `event_ts` are Aerospike Map CDT bins, not dynamic bins, which keeps the store within the 15-byte Aerospike bin-name limit regardless of how many feature views a project has. +* **Partial upserts via Map CDT ops.** Writes use `batch_write` with `map_put_items("features", {: {...}})` and `map_put("event_ts", , )`. Concurrent writes to different feature views on the same entity never clobber each other — each write mutates only its own map keys. +* **Entity-key bytes as the Aerospike user key.** Feast's `serialize_entity_key` output is passed as a `bytearray` user key (not `bytes` — the Python client hashes only the first byte of `bytes` keys, which would collapse distinct entities). +* **Timestamps as int64 epoch milliseconds.** Aerospike has no native datetime type; tz-naive timestamps are treated as UTC per the `OnlineStore` contract. + +### TTL and expiry + +`ttl_seconds` is written as record-level metadata on every `online_write_batch` call: + +| `ttl_seconds` | Aerospike TTL | Effect | +| :------------ | :-------------------------------- | :---------------------------------------------------------- | +| not set / `null` | `TTL_NAMESPACE_DEFAULT` | Record inherits the namespace's configured `default-ttl`. | +| `0` | `TTL_NEVER_EXPIRE` | Record is kept until explicitly deleted. | +| `>0` | that many seconds | Record is evicted by the server's `nsup` thread. | + +There is no per-feature-view TTL override in this version — the setting is applied uniformly for every write made by the online store. + +### Indexes + +No secondary indexes are created. All access goes through the primary key, which is the serialized entity key. + +## Async support + +Async read/write are provided by running the Aerospike Python client's blocking calls on the default thread-pool executor (`loop.run_in_executor`). The underlying C client releases the GIL during network I/O, so `await store.online_read_async(...)` keeps the event loop responsive. A native asyncio Aerospike client is not currently used. + +Both sync and async methods are fully supported: + +* `online_read` / `online_read_async` +* `online_write_batch` / `online_write_batch_async` +* `initialize` / `close` — `initialize(config)` eagerly opens the connection so feature servers pay the TCP/handshake cost at startup; `close()` releases the cached client. + +## Functionality Matrix + +The set of functionality supported by online stores is described in detail [here](overview.md#functionality). +Below is a matrix indicating which functionality is supported by the Aerospike online store. + +| | Aerospike | +| :-------------------------------------------------------- | :-------- | +| write feature values to the online store | yes | +| read feature values from the online store | yes | +| update infrastructure (e.g. tables) in the online store | yes | +| teardown infrastructure (e.g. tables) in the online store | yes | +| generate a plan of infrastructure changes | no | +| support for on-demand transforms | yes | +| readable by Python SDK | yes | +| readable by Java | no | +| readable by Go | no | +| support for entityless feature views | yes | +| support for concurrent writing to the same key | yes | +| support for ttl (time to live) at retrieval | yes | +| support for deleting expired data | yes | +| collocated by feature view | no | +| collocated by feature service | no | +| collocated by entity key | yes | + +To compare this set of functionality against other online stores, please see the full [functionality matrix](overview.md#functionality-matrix). diff --git a/docs/roadmap.md b/docs/roadmap.md index ab64289ba7b..4eac7e68b3f 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -62,6 +62,7 @@ The list below contains the functionality that contributors are planning to deve * [x] [SingleStore](https://docs.feast.dev/reference/online-stores/singlestore) * [x] [Couchbase](https://docs.feast.dev/reference/online-stores/couchbase) * [x] [MongoDB](https://docs.feast.dev/reference/online-stores/mongodb) + * [x] [Aerospike](https://docs.feast.dev/reference/online-stores/aerospike) * [x] [Qdrant (vector store)](https://docs.feast.dev/reference/online-stores/qdrant) * [x] [Milvus (vector store)](https://docs.feast.dev/reference/online-stores/milvus) * [x] [Faiss (vector store)](https://docs.feast.dev/reference/online-stores/faiss) diff --git a/examples/online_store/aerospike_overrides_and_hooks/README.md b/examples/online_store/aerospike_overrides_and_hooks/README.md new file mode 100644 index 00000000000..7143b6a3dff --- /dev/null +++ b/examples/online_store/aerospike_overrides_and_hooks/README.md @@ -0,0 +1,65 @@ +# Aerospike: per-feature-view overrides + prewriting hooks + +A short companion to [`docs/reference/online-stores/aerospike.md`](../../../docs/reference/online-stores/aerospike.md) +demonstrating three deployment patterns the Aerospike online store supports +without needing any Feast extension code: + +1. **Per-feature-view namespace overrides** — pin one view to a RAM-only + namespace and another to an SSD-backed one without splitting the project. +2. **Per-feature-view set overrides** — isolate one view in its own set so + `feast apply` deletions or admin truncates only touch that view. +3. **Prewriting hooks** — apply a project-wide write-side transformation + (PII masking in this example) without sprinkling it through every + materialization job. + +Nothing here is Aerospike-specific *infrastructure* — it's all configured +in `feature_store.yaml`. This directory only adds the hook-target Python +module the YAML references. + +## Files + +| file | purpose | +|---|---| +| [`hooks.py`](hooks.py) | A pure-Python prewriting-hook module containing `hash_pii_string_features`, the same example used in the docs. Drop into any module on the writer's `PYTHONPATH`. | +| [`feature_store.yaml`](feature_store.yaml) | Reference `online_store` block showing all three features wired together. Copy the `online_store` section into your own `feature_store.yaml` — the rest is project-specific scaffolding. | + +## Prerequisites + +* Feast installed with the Aerospike extra (`pip install 'feast[aerospike]'`). +* An Aerospike cluster reachable from your writer process. The + [Aerospike online-store reference](../../../docs/reference/online-stores/aerospike.md) + shows a minimal local CE config (`127.0.0.1:3000`); run Aerospike however + you normally would (Docker, Kubernetes, bare metal). +* On every process that calls `online_write_batch` through this store + (materialization workers, the registry CLI host, the feature server if + you run one), the `FEAST_PII_SALT` environment variable must be set + before the first write — `hash_pii_string_features` raises rather than + silently writing plaintext if the salt isn't configured. +* The two namespaces referenced by `namespace_overrides` (`feast_ram` and + `feast_ssd` in the sample YAML) must already exist on the Aerospike + cluster — Aerospike cannot create namespaces at runtime. + +## Trying it out + +1. Drop `hooks.py` into a module on your `PYTHONPATH` that the writer + process can import (e.g. inside your existing feature-repo package). + The example uses the qualified path + `examples.online_store.aerospike_overrides_and_hooks.hooks.hash_pii_string_features`. +2. Copy the `online_store:` block from `feature_store.yaml` into your + own feature repo, adjusting hosts / namespaces / the hook import path + for your project. +3. `export FEAST_PII_SALT=...` (anything random and stable across + processes — rotate by re-running materialization with a new salt). +4. `feast apply` — the new config is registered. +5. Materialize as usual — the hook runs once per `online_write_batch`, + and any feature named `email`, `phone_number` or `ssn` lands in + Aerospike as a salted SHA-256 hex digest instead of plaintext. + +## Read-side note + +Prewriting hooks are **only** invoked on the write path. If your hook is +a one-way transform (hashing, encryption-without-decryption-key) you +have to apply the same transform to the candidate value at read time +yourself. Two-way transforms (deterministic encryption, Base64) need a +matching post-read step in your serving code; the Aerospike store does +not currently expose a symmetric "postreading hook". diff --git a/examples/online_store/aerospike_overrides_and_hooks/feature_store.yaml b/examples/online_store/aerospike_overrides_and_hooks/feature_store.yaml new file mode 100644 index 00000000000..bd1061f30aa --- /dev/null +++ b/examples/online_store/aerospike_overrides_and_hooks/feature_store.yaml @@ -0,0 +1,59 @@ +# Reference feature_store.yaml demonstrating all three Aerospike +# extension points wired together. Copy the `online_store:` block into +# your own feature repo and adjust hosts / namespaces / hook import +# path for your project. +# +# Prerequisites: +# - The `feast_ram` and `feast_ssd` namespaces must already exist on +# the Aerospike cluster. Aerospike cannot create namespaces at +# runtime; a missing namespace surfaces as AEROSPIKE_ERR_PARAM on +# the first read or write touching that view. +# - `FEAST_PII_SALT` must be set in every process that calls +# online_write_batch through this store. + +project: my_feature_repo +registry: data/registry.db +provider: local + +online_store: + type: aerospike + + hosts: + - ["aerospike.internal", 3000] + + # Store-level defaults. Anything not listed in *_overrides below + # falls back to these. + namespace: feast + set_name_template: "{project}_{collection_suffix}" + + # Pin individual feature views to different namespaces -- typically + # one in-memory namespace for hot, latency-sensitive views and one + # device-backed namespace for cold, wide views. + namespace_overrides: + driver_realtime_stats: feast_ram + driver_history_lookup: feast_ssd + + # Isolate one feature view in its own set so that admin operations on + # it (truncate, scan-based deletion via `feast apply`) do not touch + # the records of other views. + set_overrides: + isolated_view: my_feature_repo_isolated + + # Project-wide write-side hook. The store dynamically imports the + # callable on first use and caches it. Adjust the import path to + # whatever module is on your writers' PYTHONPATH; the value below + # assumes you have the example folder on PYTHONPATH from the + # repository root. + prewriting_hook: examples.online_store.aerospike_overrides_and_hooks.hooks.hash_pii_string_features + + # Standard timing knobs (optional -- shown for completeness). + ttl_seconds: 86400 + read_timeout_ms: 150 + write_timeout_ms: 300 + batch_total_timeout_ms: 500 + socket_timeout_ms: 50 + max_retries: 2 + +# Offline store / entity_key_serialization_version / etc. are +# project-specific and intentionally omitted; this file is a snippet, +# not a runnable repo. diff --git a/examples/online_store/aerospike_overrides_and_hooks/hooks.py b/examples/online_store/aerospike_overrides_and_hooks/hooks.py new file mode 100644 index 00000000000..15e6f8a10c7 --- /dev/null +++ b/examples/online_store/aerospike_overrides_and_hooks/hooks.py @@ -0,0 +1,109 @@ +"""Sample prewriting hooks for the Feast Aerospike online store. + +Reference the callable from ``feature_store.yaml`` via its import string, +e.g.:: + + online_store: + type: aerospike + ... + prewriting_hook: examples.online_store.aerospike_overrides_and_hooks.hooks.hash_pii_string_features + +The Aerospike online store invokes the configured callable once per +``online_write_batch`` call, passing the rows about to be written. The +callable must return a row list with the same schema. Returning ``[]`` +short-circuits the write — same path as an empty input, no wire call is +issued. +""" + +from __future__ import annotations + +import hashlib +import os +from datetime import datetime +from typing import List, Optional, Tuple + +from feast import FeatureView +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto +from feast.repo_config import RepoConfig + +# Names of features that must never reach the online store as plaintext. +# Match is by exact feature name; tweak to your project's conventions +# (regex, suffix-based, FV-tag-driven, etc.). +_SENSITIVE_FEATURES = frozenset({"email", "phone_number", "ssn"}) + +# Type alias for the per-row payload Feast hands to ``online_write_batch``. +WriteRow = Tuple[ + EntityKeyProto, + dict, + datetime, + Optional[datetime], +] + + +def hash_pii_string_features( + config: RepoConfig, + table: FeatureView, + data: List[WriteRow], +) -> List[WriteRow]: + """Replace any sensitive string feature with a salted SHA-256 hex digest. + + Determinism: same plaintext + same ``FEAST_PII_SALT`` → same digest. + Downstream lookups that hash the candidate value the same way still + hit; lookups against the raw plaintext silently miss. + + Safety: an unset salt raises rather than falling back to plaintext. + Set ``FEAST_PII_SALT`` on every process that materialises features + (workers, registry CLI host, feature server). + """ + salt = os.environ.get("FEAST_PII_SALT") + if salt is None: + raise RuntimeError( + "FEAST_PII_SALT is not set; refusing to write feature batches " + "without a configured PII salt." + ) + salt_bytes = salt.encode("utf-8") + + def _digest(plaintext: str) -> str: + h = hashlib.sha256() + h.update(salt_bytes) + h.update(plaintext.encode("utf-8")) + return h.hexdigest() + + transformed: List[WriteRow] = [] + for entity_key, values, event_ts, created_ts in data: + new_values = dict(values) + for feature_name in _SENSITIVE_FEATURES.intersection(new_values): + v: ValueProto = new_values[feature_name] + if v.HasField("string_val") and v.string_val: + new_values[feature_name] = ValueProto(string_val=_digest(v.string_val)) + transformed.append((entity_key, new_values, event_ts, created_ts)) + return transformed + + +def drop_rows_with_negative_amounts( + config: RepoConfig, + table: FeatureView, + data: List[WriteRow], +) -> List[WriteRow]: + """Defensive sample hook: filter rows whose ``amount`` feature is < 0. + + Demonstrates that hooks can also *remove* rows. Returning an empty + list short-circuits the wire call entirely — useful for emergency + feature-write quarantines without a code deploy. + """ + keep: List[WriteRow] = [] + for entity_key, values, event_ts, created_ts in data: + amount: Optional[ValueProto] = values.get("amount") + if ( + amount is not None + and amount.HasField("double_val") + and amount.double_val < 0 + ): + continue + if amount is not None and amount.HasField("float_val") and amount.float_val < 0: + continue + if amount is not None and amount.HasField("int64_val") and amount.int64_val < 0: + continue + keep.append((entity_key, values, event_ts, created_ts)) + return keep diff --git a/infra/feast-operator/api/v1/featurestore_types.go b/infra/feast-operator/api/v1/featurestore_types.go index 830e0a77e5d..f0331ecb1b4 100644 --- a/infra/feast-operator/api/v1/featurestore_types.go +++ b/infra/feast-operator/api/v1/featurestore_types.go @@ -626,7 +626,7 @@ type OnlineStoreFilePersistence struct { // OnlineStoreDBStorePersistence configures the DB store persistence for the online store service type OnlineStoreDBStorePersistence struct { // Type of the persistence type you want to use. - // +kubebuilder:validation:Enum=snowflake.online;redis;datastore;dynamodb;bigtable;postgres;cassandra;mysql;hazelcast;singlestore;hbase;elasticsearch;qdrant;couchbase.online;milvus;hybrid;mongodb + // +kubebuilder:validation:Enum=snowflake.online;redis;datastore;dynamodb;bigtable;postgres;cassandra;mysql;hazelcast;singlestore;hbase;elasticsearch;qdrant;couchbase.online;milvus;hybrid;mongodb;aerospike Type string `json:"type"` // Data store parameters should be placed as-is from the "feature_store.yaml" under the secret key. "registry_type" & "type" fields should be removed. SecretRef corev1.LocalObjectReference `json:"secretRef"` @@ -652,6 +652,7 @@ var ValidOnlineStoreDBStorePersistenceTypes = []string{ "milvus", "hybrid", "mongodb", + "aerospike", } // LocalRegistryConfig configures the registry service diff --git a/infra/feast-operator/api/v1alpha1/featurestore_types.go b/infra/feast-operator/api/v1alpha1/featurestore_types.go index 87d13003805..c165801eda7 100644 --- a/infra/feast-operator/api/v1alpha1/featurestore_types.go +++ b/infra/feast-operator/api/v1alpha1/featurestore_types.go @@ -373,7 +373,7 @@ type OnlineStoreFilePersistence struct { // OnlineStoreDBStorePersistence configures the DB store persistence for the online store service type OnlineStoreDBStorePersistence struct { // Type of the persistence type you want to use. - // +kubebuilder:validation:Enum=snowflake.online;redis;datastore;dynamodb;bigtable;postgres;cassandra;mysql;hazelcast;singlestore;hbase;elasticsearch;qdrant;couchbase.online;milvus;hybrid;mongodb + // +kubebuilder:validation:Enum=snowflake.online;redis;datastore;dynamodb;bigtable;postgres;cassandra;mysql;hazelcast;singlestore;hbase;elasticsearch;qdrant;couchbase.online;milvus;hybrid;mongodb;aerospike Type string `json:"type"` // Data store parameters should be placed as-is from the "feature_store.yaml" under the secret key. "registry_type" & "type" fields should be removed. SecretRef corev1.LocalObjectReference `json:"secretRef"` @@ -399,6 +399,7 @@ var ValidOnlineStoreDBStorePersistenceTypes = []string{ "milvus", "hybrid", "mongodb", + "aerospike", } // LocalRegistryConfig configures the registry service diff --git a/infra/feast-operator/bundle/manifests/feast.dev_featurestores.yaml b/infra/feast-operator/bundle/manifests/feast.dev_featurestores.yaml index d8791ae8a26..1355d65d993 100644 --- a/infra/feast-operator/bundle/manifests/feast.dev_featurestores.yaml +++ b/infra/feast-operator/bundle/manifests/feast.dev_featurestores.yaml @@ -2256,6 +2256,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -8281,6 +8282,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -13578,6 +13580,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -17856,6 +17859,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef diff --git a/infra/feast-operator/config/crd/bases/feast.dev_featurestores.yaml b/infra/feast-operator/config/crd/bases/feast.dev_featurestores.yaml index 26d2d4c65f0..11650995e04 100644 --- a/infra/feast-operator/config/crd/bases/feast.dev_featurestores.yaml +++ b/infra/feast-operator/config/crd/bases/feast.dev_featurestores.yaml @@ -2399,6 +2399,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -8704,6 +8705,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -14227,6 +14229,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -18732,6 +18735,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef diff --git a/infra/feast-operator/dist/install.yaml b/infra/feast-operator/dist/install.yaml index 6c128fc320c..57853e4b83f 100644 --- a/infra/feast-operator/dist/install.yaml +++ b/infra/feast-operator/dist/install.yaml @@ -2407,6 +2407,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -8712,6 +8713,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -14235,6 +14237,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef @@ -18740,6 +18743,7 @@ spec: - milvus - hybrid - mongodb + - aerospike type: string required: - secretRef diff --git a/pixi.lock b/pixi.lock index 2e0338cfaac..4a7fdda1ef8 100644 --- a/pixi.lock +++ b/pixi.lock @@ -2421,8 +2421,8 @@ packages: requires_python: '>=3.10' - pypi: ./ name: feast - version: 0.63.1.dev113+g5c17b8d85.d20260628 - sha256: d15effc66d050583d3951905950d474615aa17f79a141ece417575d458962030 + version: 0.64.1.dev48+g10a1fb32b.d20260630 + sha256: 36b2ec53e444a094df5ce68594025cea29abcb5e164586f145d6dd7e5f92c1ff requires_dist: - click>=7.0.0,<9.0.0 - colorama>=0.3.9,<1 @@ -2455,6 +2455,7 @@ packages: - psutil - bigtree>=0.19.2 - pyjwt + - aerospike>=19.0.0,<20.0.0 ; extra == 'aerospike' - boto3>=1.38.27 ; extra == 'aws' - fsspec>=2024.1.0 ; extra == 'aws' - aiobotocore>=2 ; extra == 'aws' diff --git a/pyproject.toml b/pyproject.toml index bf2af8dc11c..506da5ea99b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,9 @@ dependencies = [ ] [project.optional-dependencies] +aerospike = [ + "aerospike>=19.0.0,<20.0.0", +] aws = ["boto3>=1.38.27", "fsspec>=2024.1.0", "aiobotocore>=2"] azure = [ "azure-storage-blob>=0.37.0", diff --git a/sdk/python/feast/infra/online_stores/aerospike_online_store/__init__.py b/sdk/python/feast/infra/online_stores/aerospike_online_store/__init__.py new file mode 100644 index 00000000000..099892626dd --- /dev/null +++ b/sdk/python/feast/infra/online_stores/aerospike_online_store/__init__.py @@ -0,0 +1,3 @@ +from .aerospike import AerospikeOnlineStore, AerospikeOnlineStoreConfig + +__all__ = ["AerospikeOnlineStore", "AerospikeOnlineStoreConfig"] diff --git a/sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike.py b/sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike.py new file mode 100644 index 00000000000..208c3b0684a --- /dev/null +++ b/sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike.py @@ -0,0 +1,1047 @@ +from __future__ import annotations + +import asyncio +import functools +import importlib +import threading +from datetime import datetime, timezone +from logging import getLogger +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + Literal, + Optional, + Sequence, + Set, + Tuple, + TypeVar, + Union, +) + +from pydantic import SecretStr + +try: + import aerospike + from aerospike_helpers import cdt_ctx + from aerospike_helpers.batch.records import BatchRecords + from aerospike_helpers.batch.records import Write as BatchWrite + from aerospike_helpers.operations import map_operations as map_ops + from aerospike_helpers.operations import operations as ops +except ImportError as e: + from feast.errors import FeastExtrasDependencyImportError + + raise FeastExtrasDependencyImportError("aerospike", str(e)) + +from feast.entity import Entity +from feast.feature_view import FeatureView +from feast.infra.key_encoding_utils import serialize_entity_key +from feast.infra.online_stores.online_store import OnlineStore +from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto +from feast.protos.feast.types.Value_pb2 import Value as ValueProto +from feast.repo_config import FeastConfigBaseModel, RepoConfig +from feast.type_map import ( + feast_value_type_to_python_type, + python_values_to_proto_values, +) + +logger = getLogger(__name__) + + +# Aerospike per-record batch result codes we treat specially. Anything not +# listed here is surfaced as an exception so a transient server error (e.g. +# timeout, device overload) is never silently misreported as a missing feature. +# See https://aerospike.com/docs/server/reference/errors. +_AS_OK: int = 0 +_AS_ERR_RECORD_NOT_FOUND: int = 2 # genuine "entity has never been written" +_AS_ERR_OP_NOT_APPLICABLE: int = 26 # map/list op targeted a missing CDT path + + +_AUTH_MODE_TO_CONSTANT: Dict[str, int] = { + "internal": aerospike.AUTH_INTERNAL, + "external": aerospike.AUTH_EXTERNAL, + "pki": aerospike.AUTH_PKI, +} + + +# Type alias for the prewriting-hook signature. Hooks receive the rows about +# to be written and return a (possibly transformed) row list with the same +# schema. Defined as ``Any`` in the value position to keep the public type +# hint readable; the actual contract is documented on +# ``AerospikeOnlineStoreConfig.prewriting_hook``. +PrewritingHook = Callable[ + [ + "RepoConfig", + "FeatureView", + List[ + Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]] + ], + ], + List[Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]], +] + + +# Create every Map CDT bin with an ordered map. map_get_by_key / +# map_remove_by_key on an ordered map are O(log N) in the map size instead of +# O(N), which matters on the update() background scan (which walks every +# record in the project's set) and on reads of wide feature views. The policy +# is applied on each put so map-creation on the first write picks up the +# ordering; subsequent puts keep it. +_ORDERED_MAP_POLICY: Dict[str, Any] = {"map_order": aerospike.MAP_KEY_ORDERED} + +# Aerospike server ``batch-max-requests`` defaults to 5000 on many clusters (0 +# means unlimited on newer releases). Stay well under that so materialization +# and wide feature-server requests do not trip BatchMaxRequestError (code 151). +_DEFAULT_BATCH_MAX_RECORDS: int = 1_000 + +_T = TypeVar("_T") + + +def _datetime_to_epoch_ms(dt: datetime) -> int: + """Convert a datetime to int64 epoch milliseconds. + + Aerospike has no native datetime type, so timestamps are stored as int + bins. Per the OnlineStore contract, a tz-naive timestamp is treated as UTC. + """ + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return int(dt.timestamp() * 1000) + + +def _epoch_ms_to_datetime(value: Optional[int]) -> Optional[datetime]: + """Inverse of :func:`_datetime_to_epoch_ms`. Returns a tz-aware UTC datetime.""" + if value is None: + return None + return datetime.fromtimestamp(value / 1000.0, tz=timezone.utc) + + +def _resolve_ttl(ttl_seconds: Optional[int]) -> int: + """Map the config's ``ttl_seconds`` to an Aerospike record-metadata TTL. + + * ``None`` -> use the namespace default (``TTL_NAMESPACE_DEFAULT``). + * ``0`` -> never expire (``TTL_NEVER_EXPIRE``). + * ``> 0`` -> that many seconds until expiry. + """ + if ttl_seconds is None: + return aerospike.TTL_NAMESPACE_DEFAULT + if ttl_seconds == 0: + return aerospike.TTL_NEVER_EXPIRE + return int(ttl_seconds) + + +class AerospikeOnlineStoreConfig(FeastConfigBaseModel): + """Aerospike configuration. + + Aerospike does not have a URI analogue; connections are established via a + seed list of ``(host, port)`` or ``(host, port, tls_name)`` tuples. See the + Aerospike Python client reference for the meaning of additional policies and + TLS options surfaced below, and use ``client_kwargs`` for anything not + explicitly modelled here. + """ + + type: Literal["aerospike"] = "aerospike" + """Online store type selector""" + + hosts: List[Union[Tuple[str, int], Tuple[str, int, str]]] = [("localhost", 3000)] + """Aerospike seed nodes. + + Each entry is either ``(host, port)`` or ``(host, port, tls_name)`` when TLS + is enabled. At least one seed node is required. + """ + + namespace: str = "feast" + """Default Aerospike namespace. Must be pre-configured on the cluster — + namespaces cannot be created at runtime. Used for any feature view not + listed in :attr:`namespace_overrides`.""" + + set_name_template: str = "{project}_{collection_suffix}" + """Template for the per-project Aerospike set name. Available substitutions: + ``{project}`` and ``{collection_suffix}``. Used for any feature view not + listed in :attr:`set_overrides`.""" + + collection_suffix: str = "latest" + """Suffix used by ``set_name_template`` to distinguish sets belonging to the + same project (e.g. a future multi-version layout).""" + + namespace_overrides: Dict[str, str] = {} + """Per-feature-view namespace overrides. Maps a feature view name to the + Aerospike namespace that view should be stored in. Falls back to + :attr:`namespace` for any feature view not listed. + + Lets a deployment pin hot small views to a RAM-only namespace and wider + archival views to an SSD-backed one without having to split projects:: + + namespace_overrides = { + "driver_realtime_stats": "feast_ram", + "driver_history_lookup": "feast_ssd", + } + + Every namespace listed here MUST already exist on the cluster — Aerospike + cannot create namespaces at runtime, and a missing namespace surfaces as + an opaque ``AEROSPIKE_ERR_PARAM`` on the first read or write. + """ + + set_overrides: Dict[str, str] = {} + """Per-feature-view set name overrides. Maps a feature view name to a + fully-qualified Aerospike set name. Falls back to the rendering of + :attr:`set_name_template` for any feature view not listed. + + Useful when a deployment wants each feature view in its own set so admin + operations like ``truncate`` or the ``feast apply`` deletion path can + target one view without scanning the records of the others. Tradeoff: + multi-feature-view reads on the same entity become multiple round trips + instead of one — only set this when the operational isolation is worth + that cost. + """ + + prewriting_hook: Optional[str] = None + """Optional import path of a callable applied to every batch of rows + before they are written. Used for cross-cutting concerns like PII + masking, encryption-at-rest, value coercion or dual-write fan-out. + + Format: ``"package.module.function_name"``. Resolved by import string + (rather than a Python ``Callable``) so the config survives YAML / JSON + serialisation and remote-feature-server transport. + + The hook signature is :data:`PrewritingHook`:: + + def hook( + config: RepoConfig, + table: FeatureView, + data: List[Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]], + ) -> List[Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]]: + ... + + The hook MUST return a row list with the same shape as its input. + Hooks that raise will fail the whole batch — there is no per-row + fallback. The resolved callable is cached on the store instance; if the + configured import string changes between calls, it is re-resolved + automatically on the next write. + """ + + user: Optional[str] = None + """Optional username for Aerospike Enterprise authentication.""" + + password: Optional[SecretStr] = None + """Optional password for Aerospike Enterprise authentication.""" + + auth_mode: Literal["internal", "external", "pki"] = "internal" + """Authentication mode. ``internal`` for CE/EE user/password, ``external`` + for LDAP/Kerberos, ``pki`` for certificate-based auth.""" + + tls: Optional[Dict[str, Any]] = None + """TLS configuration, passed through verbatim to the Aerospike client. + See the Aerospike Python client ``tls`` policy options.""" + + ttl_seconds: Optional[int] = None + """Record-level TTL, applied to every write. ``None`` uses the namespace + default, ``0`` means never expire (mapped to the client's ``-1`` sentinel). + No per-feature-view override in v1.""" + + write_timeout_ms: int = 1_000 + """Per-call write total timeout in milliseconds. This is the hard deadline + the client gives a single ``put`` / ``operate`` — including any retries — + to return a response, after which the call fails.""" + + read_timeout_ms: int = 250 + """Per-call read total timeout in milliseconds. Hard deadline for a + single-record ``get`` — including any retries — after which the call + fails.""" + + batch_total_timeout_ms: int = 2_000 + """Total timeout in milliseconds for a whole ``batch_write`` / + ``batch_operate`` call, including retries. Applies to every batch + operation ``online_read`` and ``online_write_batch`` issue.""" + + socket_timeout_ms: Optional[int] = None + """Per-attempt socket timeout in milliseconds. This is the per-retry + trigger that lets ``max_retries`` actually fire within the caller's + overall ``*_timeout_ms`` budget — without it, a single attempt can + consume the whole deadline and retries never run. Applied uniformly + to ``read``, ``write`` and ``batch`` policies. ``None`` leaves the + client default in place.""" + + max_retries: int = 2 + """Maximum number of automatic retries on transient errors.""" + + batch_max_records: int = _DEFAULT_BATCH_MAX_RECORDS + """Maximum records per ``batch_write`` / ``batch_operate`` call. + + Aerospike enforces a per-node batch size via the server ``batch-max-requests`` + setting (historically 5000). Feast chunks read and write paths to this limit + so large materializations and wide online-serving requests do not fail the + whole batch when the server cap is exceeded. + """ + + client_kwargs: Dict[str, Any] = {} + """Escape hatch for any Aerospike client configuration not surfaced above. + Merged into the client config passed to ``aerospike.client()``.""" + + +class AerospikeOnlineStore(OnlineStore): + """Aerospike implementation of the Feast :class:`OnlineStore`. + + Storage layout (one set per project): + + * Namespace: ``config.online_store.namespace`` (server-configured) + * Set: ``{project}_{collection_suffix}`` + * Key: ``serialize_entity_key(entity_key)`` (bytes) + * Bins: + + * ``features`` — Map CDT ``{"": {"": }}`` + * ``event_ts`` — Map CDT ``{"": }`` + * ``created_ts`` — top-level ```` + + Timestamps are stored as int64 epoch milliseconds because Aerospike has no + native datetime type. Tz-naive timestamps are treated as UTC per the + :class:`OnlineStore` contract. + """ + + def __init__(self) -> None: + # Kept on the instance rather than the class so two ``AerospikeOnlineStore`` + # instances can't accidentally share a cached client through class state. + self._client: Optional[aerospike.Client] = None + # Resolved prewriting hook, cached on the instance so the import + # cost is paid once. ``_prewriting_hook_spec`` records the import + # string the cache was built from so a config swap (re-binding the + # store to a different hook) re-resolves on next call. + self._prewriting_hook: Optional[PrewritingHook] = None + self._prewriting_hook_spec: Optional[str] = None + self._client_lock = threading.Lock() + + # ------------------------------------------------------------------ + # Lifecycle / connection management + # ------------------------------------------------------------------ + @staticmethod + def _chunked(items: Sequence[_T], chunk_size: int) -> Iterator[Sequence[_T]]: + """Yield slices of ``items`` no larger than ``chunk_size``.""" + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + for start in range(0, len(items), chunk_size): + yield items[start : start + chunk_size] + + def _batch_max_records(self, config: RepoConfig) -> int: + store_cfg = config.online_store + if not isinstance(store_cfg, AerospikeOnlineStoreConfig): + raise RuntimeError(f"{config.online_store.type = }. It must be aerospike.") + return store_cfg.batch_max_records + + def _get_client(self, config: RepoConfig) -> aerospike.Client: + """Lazily create and cache an Aerospike client on first use. + + The underlying C client maintains its own connection pool, so a single + cached instance is safe to share across calls on this store. Creation + is guarded by a lock so concurrent first callers in a threaded feature + server do not leak extra connections. + """ + if self._client is not None: + return self._client + + with self._client_lock: + if self._client is not None: + return self._client + + if not isinstance(config.online_store, AerospikeOnlineStoreConfig): + raise RuntimeError( + f"{config.online_store.type = }. It must be aerospike." + ) + store_cfg = config.online_store + + read_policy: Dict[str, Any] = { + "total_timeout": store_cfg.read_timeout_ms, + "max_retries": store_cfg.max_retries, + } + write_policy: Dict[str, Any] = { + "total_timeout": store_cfg.write_timeout_ms, + "max_retries": store_cfg.max_retries, + } + batch_policy: Dict[str, Any] = { + "total_timeout": store_cfg.batch_total_timeout_ms, + "max_retries": store_cfg.max_retries, + } + if store_cfg.socket_timeout_ms is not None: + # socket_timeout is the per-attempt deadline; without it, + # total_timeout is the whole budget and retries never fire. + read_policy["socket_timeout"] = store_cfg.socket_timeout_ms + write_policy["socket_timeout"] = store_cfg.socket_timeout_ms + batch_policy["socket_timeout"] = store_cfg.socket_timeout_ms + + client_config: Dict[str, Any] = { + "hosts": [tuple(h) for h in store_cfg.hosts], + "policies": { + "read": read_policy, + "write": write_policy, + "batch": batch_policy, + }, + **store_cfg.client_kwargs, + } + if store_cfg.user: + if store_cfg.password is None: + raise ValueError( + "AerospikeOnlineStoreConfig.user is set but password is not." + ) + client_config["user"] = store_cfg.user + client_config["password"] = store_cfg.password.get_secret_value() + client_config["auth_mode"] = _AUTH_MODE_TO_CONSTANT[store_cfg.auth_mode] + if store_cfg.tls: + client_config["tls"] = store_cfg.tls + + self._client = aerospike.client(client_config).connect() + return self._client + + def _set_name(self, config: RepoConfig, fv_name: Optional[str] = None) -> str: + """Resolve the Aerospike set name for a feature view. + + Without ``fv_name`` returns the project-level default rendered from + ``set_name_template``. With an ``fv_name`` returns the override from + :attr:`AerospikeOnlineStoreConfig.set_overrides` if present, falling + back to the project default. Calling without ``fv_name`` is preserved + for callers (and tests) that want the project's default set without + knowing which feature views might override it. + """ + store_cfg = config.online_store + overrides = getattr(store_cfg, "set_overrides", None) or {} + if fv_name is not None and fv_name in overrides: + return overrides[fv_name] + return store_cfg.set_name_template.format( + project=config.project, + collection_suffix=store_cfg.collection_suffix, + ) + + def _namespace_for_fv( + self, config: RepoConfig, fv_name: Optional[str] = None + ) -> str: + """Resolve the Aerospike namespace for a feature view. + + Mirrors :meth:`_set_name` semantics: without ``fv_name`` returns the + store-default namespace; with one, returns the override from + :attr:`AerospikeOnlineStoreConfig.namespace_overrides` if present. + """ + store_cfg = config.online_store + overrides = getattr(store_cfg, "namespace_overrides", None) or {} + if fv_name is not None and fv_name in overrides: + return overrides[fv_name] + return store_cfg.namespace + + def _aerospike_key( + self, + config: RepoConfig, + entity_key: EntityKeyProto, + fv_name: Optional[str] = None, + ) -> Tuple[str, str, bytearray]: + """Build a ``(namespace, set, user_key)`` tuple for an entity. + + When ``fv_name`` is provided, the namespace and set name honour the + per-feature-view overrides from + :class:`AerospikeOnlineStoreConfig`. Without ``fv_name`` the + store-level defaults are used. + + The user key is returned as a ``bytearray`` rather than ``bytes``: + the Aerospike Python C client rejects ``bytes`` user keys + (``calc_digest`` raises ``"Key is invalid"``), and ``batch_read`` / + ``batch_operate`` silently hash only the first byte of a ``bytes`` + key. ``bytearray`` is the supported binary-key type. + """ + user_key = serialize_entity_key( + entity_key, + entity_key_serialization_version=config.entity_key_serialization_version, + ) + return ( + self._namespace_for_fv(config, fv_name), + self._set_name(config, fv_name), + bytearray(user_key), + ) + + def _resolve_prewriting_hook(self, config: RepoConfig) -> Optional[PrewritingHook]: + """Resolve and cache the configured prewriting hook, if any. + + The hook is referenced by import string in the config (so it + survives YAML/JSON round-trips and remote-feature-server transport), + and resolved here on first use. The resolved callable is cached on + the store instance — and re-resolved automatically if a subsequent + config swap rebinds the store to a different hook. + """ + spec = getattr(config.online_store, "prewriting_hook", None) + if not spec: + # A previously-set hook on this instance must be cleared if the + # caller has since unset it on the config; otherwise we'd keep + # applying a hook the user explicitly turned off. + self._prewriting_hook = None + self._prewriting_hook_spec = None + return None + if spec == self._prewriting_hook_spec and self._prewriting_hook is not None: + return self._prewriting_hook + + module_path, _, fn_name = spec.rpartition(".") + if not module_path or not fn_name: + raise ValueError( + "AerospikeOnlineStoreConfig.prewriting_hook must be a fully " + f"qualified import path 'package.module.function', got: {spec!r}" + ) + try: + module = importlib.import_module(module_path) + except ImportError as e: + raise ValueError( + f"prewriting_hook {spec!r}: could not import module " + f"{module_path!r}: {e}" + ) from e + try: + hook = getattr(module, fn_name) + except AttributeError: + raise ValueError( + f"prewriting_hook {spec!r}: module {module_path!r} has no " + f"attribute {fn_name!r}" + ) from None + if not callable(hook): + raise TypeError( + f"prewriting_hook {spec!r} resolved to a non-callable " + f"{type(hook).__name__}" + ) + self._prewriting_hook = hook + self._prewriting_hook_spec = spec + return hook + + # ------------------------------------------------------------------ + # Write path + # ------------------------------------------------------------------ + @staticmethod + def _build_batch_writes( + config: RepoConfig, + table: FeatureView, + data: List[ + Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]] + ], + namespace: str, + set_name: str, + ) -> BatchRecords: + """Build a :class:`BatchRecords` with one :class:`BatchWrite` per row. + + Each row becomes an atomic, server-side Map-put op list: + + * ``features[][] = `` for every requested feature + (single ``map_put_items`` op keyed by feature-view name). + * ``event_ts[] = ``. + * ``created_ts = `` when provided. + + Using Map CDT ops rather than a full-record ``put`` means two writers + touching different feature views on the same entity will not clobber + each other — each write only mutates its own slot in the outer map. + + The Map CDTs are created with ``MAP_KEY_ORDERED`` so key lookups on + reads and the ``update()`` background scan stay O(log N) in the map + size. Writes use the default ``POLICY_KEY_DIGEST`` — the serialized + entity key itself is not stored on the server, saving per-record + storage that the read path never consumes (result order is preserved + by ``batch_operate`` and paired back via ``zip`` in ``online_read``). + """ + ttl_meta = {"ttl": _resolve_ttl(config.online_store.ttl_seconds)} + + batch = BatchRecords() + for entity_key, proto_values, event_timestamp, created_timestamp in data: + user_key = bytearray( + serialize_entity_key( + entity_key, + entity_key_serialization_version=config.entity_key_serialization_version, + ) + ) + feature_map = { + field: feast_value_type_to_python_type(val) + for field, val in proto_values.items() + } + operations: List[Dict[str, Any]] = [ + map_ops.map_put_items( + "features", + {table.name: feature_map}, + map_policy=_ORDERED_MAP_POLICY, + ), + map_ops.map_put( + "event_ts", + table.name, + _datetime_to_epoch_ms(event_timestamp), + map_policy=_ORDERED_MAP_POLICY, + ), + ] + if created_timestamp is not None: + operations.append( + ops.write("created_ts", _datetime_to_epoch_ms(created_timestamp)) + ) + batch.batch_records.append( + BatchWrite( + key=(namespace, set_name, user_key), + ops=operations, + meta=ttl_meta, + ) + ) + return batch + + def online_write_batch( + self, + config: RepoConfig, + table: FeatureView, + data: List[ + Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]] + ], + progress: Optional[Callable[[int], Any]], + ) -> None: + """Write a batch of feature rows using Aerospike's native batch-write API. + + Each row is upserted as a set of Map CDT operations on a single record, + preserving data for other feature views that share the same entity key. + """ + if not data: + if progress: + progress(0) + return + + # Hook resolution happens before the early-data check to avoid wiring + # work, but after the empty-batch check so a "no rows" call doesn't + # pay the import cost. Hooks are allowed to mutate ``data`` in place + # or return a new list; we always rebind to whatever they return. + hook = self._resolve_prewriting_hook(config) + if hook is not None: + data = hook(config, table, data) + if not data: + if progress: + progress(0) + return + + client = self._get_client(config) + namespace = self._namespace_for_fv(config, table.name) + set_name = self._set_name(config, table.name) + chunk_size = self._batch_max_records(config) + written = 0 + for chunk in self._chunked(data, chunk_size): + batch = self._build_batch_writes( + config, table, list(chunk), namespace, set_name + ) + if batch.batch_records: + client.batch_write(batch) + # Per-record result codes must be inspected: client.batch_write + # only raises if the whole request was rejected. A partial failure + # (e.g. a single-partition timeout) is otherwise silent, which in + # an online-serving path presents downstream as "model saw stale + # features" weeks after the fact. + self._raise_on_batch_errors(batch.batch_records, set_name, op="write") + written += len(chunk) + if progress: + progress(written) + + # ------------------------------------------------------------------ + # Read path + # ------------------------------------------------------------------ + def online_read( + self, + config: RepoConfig, + table: FeatureView, + entity_keys: List[EntityKeyProto], + requested_features: Optional[List[str]] = None, + ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]: + """Read feature values for a batch of entities in a single round trip. + + Uses Aerospike's ``batch_operate`` with two server-side Map-get ops per + record. When ``requested_features`` is provided, only those feature + columns are shipped over the wire using ``map_get_by_key_list`` nested + into the feature-view's Map CDT via ``cdt_ctx_map_key``. Otherwise the + whole feature-view slot is returned. + + * features op (projected): + ``map_get_by_key_list("features", requested_features, + MAP_RETURN_KEY_VALUE, ctx=[cdt_ctx_map_key()])`` + * features op (full): + ``map_get_by_key("features", , MAP_RETURN_VALUE)`` + * event_ts op (always): + ``map_get_by_key("event_ts", , MAP_RETURN_VALUE)`` + + Per-record status codes are inspected so we can tell a genuine miss + (``RECORD_NOT_FOUND``, or a nested ``OP_NOT_APPLICABLE`` when the + feature-view slot is absent) apart from a transient server error, + which is raised rather than silently returned as a null row. Output + order matches ``entity_keys``. + """ + if not entity_keys: + return [] + + client = self._get_client(config) + ns = self._namespace_for_fv(config, table.name) + set_name = self._set_name(config, table.name) + read_ops = self._build_read_ops(table.name, requested_features) + chunk_size = self._batch_max_records(config) + + # ``ids`` and ``docs`` use immutable ``bytes`` because ``bytearray`` is + # unhashable and can't key a dict. Keys on the wire must stay + # ``bytearray`` (see ``_aerospike_key``) — we only convert here for + # lookup. + ids: List[bytes] = [] + docs: Dict[bytes, Dict[str, Any]] = {} + + for entity_chunk in self._chunked(entity_keys, chunk_size): + keys = [ + ( + ns, + set_name, + bytearray( + serialize_entity_key( + k, + entity_key_serialization_version=config.entity_key_serialization_version, + ) + ), + ) + for k in entity_chunk + ] + batch = client.batch_operate(keys, read_ops) + chunk_ids = [bytes(user_key) for _, _, user_key in keys] + ids.extend(chunk_ids) + # batch_operate preserves request order. We pair each response with + # the original user-key rather than ``br.key[2]``: the Aerospike + # client may return the key in a different representation (e.g. only + # the first byte as a str when the write didn't use POLICY_KEY_SEND + # for reads). + for user_key, br in zip(chunk_ids, batch.batch_records): + if br.result == _AS_ERR_RECORD_NOT_FOUND: + continue + if br.result == _AS_ERR_OP_NOT_APPLICABLE: + # The record exists but the nested feature-view slot doesn't; + # treat as a miss to match the OnlineStore contract. + continue + if br.result != _AS_OK: + raise RuntimeError( + f"Aerospike batch_operate returned a non-OK status for " + f"entity (ns={ns}, set={set_name}): result={br.result}" + ) + if br.record is None: + continue + _, _, bins = br.record + raw_features = bins.get("features") if bins else None + fv_event_ts_ms = bins.get("event_ts") if bins else None + fv_features = self._normalize_projected_features(raw_features) + docs[user_key] = { + "features": {table.name: fv_features} + if fv_features is not None + else {}, + "event_timestamps": { + table.name: _epoch_ms_to_datetime(fv_event_ts_ms) + }, + } + + return self._convert_raw_docs_to_proto(ids, docs, table) + + @staticmethod + def _build_read_ops( + fv_name: str, requested_features: Optional[List[str]] + ) -> List[Dict[str, Any]]: + """Build the per-record op list for an ``online_read`` call. + + Projects ``requested_features`` server-side via + ``map_get_by_key_list`` + ``cdt_ctx_map_key`` when a projection list + is provided. Without a projection, returns the whole feature-view + submap. + """ + if requested_features: + features_op = map_ops.map_get_by_key_list( + "features", + list(requested_features), + aerospike.MAP_RETURN_KEY_VALUE, + ctx=[cdt_ctx.cdt_ctx_map_key(fv_name)], + ) + else: + features_op = map_ops.map_get_by_key( + "features", fv_name, aerospike.MAP_RETURN_VALUE + ) + return [ + features_op, + map_ops.map_get_by_key("event_ts", fv_name, aerospike.MAP_RETURN_VALUE), + ] + + @staticmethod + def _normalize_projected_features( + raw: Optional[Union[Dict[str, Any], List[Any]]], + ) -> Optional[Dict[str, Any]]: + """Convert an Aerospike features payload into a uniform ``{name: val}`` dict. + + The shape depends on which op produced the payload: + + * ``map_get_by_key("features", , MAP_RETURN_VALUE)`` returns a + ``dict`` (the inner feature-view submap). + * ``map_get_by_key_list("features", [...], MAP_RETURN_KEY_VALUE, + ctx=...)`` returns a flat ``list`` of ``[k1, v1, k2, v2, ...]`` + containing only the requested keys that exist. + """ + if raw is None: + return None + if isinstance(raw, dict): + return raw + if isinstance(raw, list): + if not raw: + return {} + return dict(zip(raw[0::2], raw[1::2])) + return None + + @staticmethod + def _raise_on_batch_errors( + batch_records: Sequence[Any], set_name: str, op: str + ) -> None: + """Raise if any per-record result code signals a failed batch write/op. + + ``client.batch_write`` and ``client.batch_operate`` only raise when the + overall request was rejected; partial failures (a single-partition + timeout, a replica quorum miss, etc.) are surfaced per record via + ``br.result`` and are otherwise silent. In an online-serving path + those silent failures later present as missing features, so we fail + loud here instead. + """ + errors = [br.result for br in batch_records if br.result != _AS_OK] + if errors: + raise RuntimeError( + f"Aerospike batch_{op} returned non-OK status codes for " + f"{len(errors)} of {len(batch_records)} records " + f"(set={set_name}): codes={errors[:10]}" + ) + + @staticmethod + def _convert_raw_docs_to_proto( + ids: List[bytes], + docs: Dict[bytes, Dict[str, Any]], + table: FeatureView, + ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]: + """Convert raw feature maps into ordered proto rows. + + The heavy lifting is done by + :func:`feast.type_map.python_values_to_proto_values`, which is + column-oriented and expects a list of values of a single type. This + helper transforms the row-oriented Aerospike lookup result into + columns, converts each column once, then reassembles rows — mirroring + the MongoDB online store's reshape so we amortize the python→proto + cost across the whole batch. + + Args: + ids: serialized entity-key bytes, in the order requested. + docs: ``{entity_id_bytes: {"features": {: {...}}, + "event_timestamps": {: datetime}}}``. Missing keys denote + "record not found". + table: FeatureView being read; provides feature name → type. + + Returns: + A list of ``(event_timestamp, feature_dict)`` the same length as + ``ids`` (``(None, None)`` for entities that had no data for this + feature view). + """ + feature_type_map = { + feature.name: feature.dtype.to_value_type() for feature in table.features + } + + raw_feature_columns: Dict[str, List[Any]] = { + feature_name: [] for feature_name in feature_type_map + } + for entity_id in ids: + doc = docs.get(entity_id) + feature_dict = doc.get("features", {}).get(table.name, {}) if doc else {} + for feature_name in feature_type_map: + raw_feature_columns[feature_name].append( + feature_dict.get(feature_name, None) + ) + + proto_feature_columns = { + feature_name: python_values_to_proto_values( + raw_values, feature_type=feature_type_map[feature_name] + ) + for feature_name, raw_values in raw_feature_columns.items() + } + + results: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = [] + for i, entity_id in enumerate(ids): + doc = docs.get(entity_id) + if doc is None: + results.append((None, None)) + continue + + fv_features = doc.get("features", {}).get(table.name) + if fv_features is None: + results.append((None, None)) + continue + + ts = doc.get("event_timestamps", {}).get(table.name) + row_features = { + feature_name: proto_feature_columns[feature_name][i] + for feature_name in proto_feature_columns + } + results.append((ts, row_features)) + return results + + # ------------------------------------------------------------------ + # Async wrappers + # ------------------------------------------------------------------ + # The Aerospike Python client is a synchronous C extension; there is no + # native asyncio interface. Network calls do release the GIL, so we expose + # a correct ``async`` surface by offloading each blocking call to the + # default thread-pool executor. Callers that ``await`` these methods keep + # the event loop responsive while the client talks to the cluster. + + async def online_write_batch_async( + self, + config: RepoConfig, + table: FeatureView, + data: List[ + Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]] + ], + progress: Optional[Callable[[int], Any]], + ) -> None: + loop = asyncio.get_running_loop() + await loop.run_in_executor( + None, + functools.partial(self.online_write_batch, config, table, data, progress), + ) + + async def online_read_async( + self, + config: RepoConfig, + table: FeatureView, + entity_keys: List[EntityKeyProto], + requested_features: Optional[List[str]] = None, + ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, + functools.partial( + self.online_read, config, table, entity_keys, requested_features + ), + ) + + async def initialize(self, config: RepoConfig) -> None: + """Pre-warm the Aerospike client so the first request is hot. + + Feature servers typically call :meth:`initialize` during startup so the + TCP + handshake latency is paid upfront rather than on the first + ``online_read``. + """ + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, functools.partial(self._get_client, config)) + + async def close(self) -> None: + """Release the cached Aerospike client, if any.""" + with self._client_lock: + if self._client is None: + return + client = self._client + self._client = None + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, client.close) + + # ------------------------------------------------------------------ + # Admin paths (update / teardown) + # ------------------------------------------------------------------ + def update( + self, + config: RepoConfig, + tables_to_delete: Sequence[FeatureView], + tables_to_keep: Sequence[FeatureView], + entities_to_delete: Sequence[Entity], + entities_to_keep: Sequence[Entity], + partial: bool, + ) -> None: + """Reconcile per-feature-view data when a schema change is applied. + + Aerospike has no explicit schema, and records/sets are created lazily + on first write, so there is nothing to do for ``tables_to_keep`` or + either of the entity lists. For ``tables_to_delete`` we strip each + feature-view's slot out of the ``features`` and ``event_ts`` Map + CDTs on every record that contains it. + + Dropped feature views are grouped by their resolved + ``(namespace, set)`` (per :attr:`namespace_overrides` / + :attr:`set_overrides`), and each group is issued as one + **background scan** with a combined op list — a single server-side + pass per (ns, set), regardless of how many feature views in that + group are being dropped. Without this grouping, a deployment that + spreads feature views across multiple namespaces or sets would + either issue blind cross-namespace scans (impossible) or scan one + per dropped FV (wasteful). The scans run asynchronously server-side + and return immediately; this matches the intent of ``feast apply``, + after which the caller stops reading the dropped feature views + anyway. + """ + if not isinstance(config.online_store, AerospikeOnlineStoreConfig): + raise RuntimeError(f"{config.online_store.type = }. It must be aerospike.") + if not tables_to_delete: + return + + client = self._get_client(config) + + # Group dropped feature views by (namespace, set) so each scan only + # touches the records that could plausibly contain the dropped slot. + groups: Dict[Tuple[str, str], List[FeatureView]] = {} + for fv in tables_to_delete: + key = ( + self._namespace_for_fv(config, fv.name), + self._set_name(config, fv.name), + ) + groups.setdefault(key, []).append(fv) + + for (ns, set_name), fvs in groups.items(): + remove_ops: List[Dict[str, Any]] = [] + for fv in fvs: + remove_ops.append( + map_ops.map_remove_by_key( + "features", fv.name, aerospike.MAP_RETURN_NONE + ) + ) + remove_ops.append( + map_ops.map_remove_by_key( + "event_ts", fv.name, aerospike.MAP_RETURN_NONE + ) + ) + + scan = client.scan(ns, set_name) + scan.add_ops(remove_ops) + scan.execute_background() + + def teardown( + self, + config: RepoConfig, + tables: Sequence[FeatureView], + entities: Sequence[Entity], + ) -> None: + """Truncate every (namespace, set) the project may have written to + and close the cached client. + + Uses Aerospike's ``truncate(namespace, set, 0)`` — a set-scoped + metadata operation that clears every record in O(1) client time, + cheaper than Mongo's ``collection.drop()``. Passing ``0`` as the + cutoff means "drop everything regardless of last-update time". + + Collects the unique ``(namespace, set)`` pairs from the project's + store-level default plus every feature view in ``tables`` (each + resolved through :attr:`namespace_overrides` / + :attr:`set_overrides`). The default is always included so a + teardown invoked with an empty ``tables`` list still clears the + store-default location. + + Truncate on a non-existent set is a no-op, so calling ``teardown`` + on a project that never wrote data — or on a feature view that + was never written to — is safe. + """ + if not isinstance(config.online_store, AerospikeOnlineStoreConfig): + raise RuntimeError(f"{config.online_store.type = }. It must be aerospike.") + + client = self._get_client(config) + + pairs: Set[Tuple[str, str]] = { + (config.online_store.namespace, self._set_name(config)) + } + for fv in tables: + pairs.add( + ( + self._namespace_for_fv(config, fv.name), + self._set_name(config, fv.name), + ) + ) + # sort to keep the truncation order deterministic for tests; the + # operation is independent per (ns, set) so order has no semantic + # meaning. + for ns, set_name in sorted(pairs): + client.truncate(ns, set_name, 0) + with self._client_lock: + if self._client is not None: + self._client.close() + self._client = None diff --git a/sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike_repo_configuration.py b/sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike_repo_configuration.py new file mode 100644 index 00000000000..e704aca780d --- /dev/null +++ b/sdk/python/feast/infra/online_stores/aerospike_online_store/aerospike_repo_configuration.py @@ -0,0 +1,13 @@ +from tests.universal.feature_repos.integration_test_repo_config import ( + IntegrationTestRepoConfig, +) +from tests.universal.feature_repos.universal.online_store.aerospike import ( + AerospikeOnlineStoreCreator, +) + +FULL_REPO_CONFIGS = [ + IntegrationTestRepoConfig( + online_store="aerospike", + online_store_creator=AerospikeOnlineStoreCreator, + ), +] diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index 56401b4f6f1..847f933d4f6 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -66,6 +66,7 @@ } ONLINE_STORE_CLASS_FOR_TYPE = { + "aerospike": "feast.infra.online_stores.aerospike_online_store.AerospikeOnlineStore", "sqlite": "feast.infra.online_stores.sqlite.SqliteOnlineStore", "datastore": "feast.infra.online_stores.datastore.DatastoreOnlineStore", "redis": "feast.infra.online_stores.redis.RedisOnlineStore", diff --git a/sdk/python/tests/unit/online_store/test_aerospike_online_retrieval.py b/sdk/python/tests/unit/online_store/test_aerospike_online_retrieval.py new file mode 100644 index 00000000000..f44d27bb0d7 --- /dev/null +++ b/sdk/python/tests/unit/online_store/test_aerospike_online_retrieval.py @@ -0,0 +1,1684 @@ +""" +Unit tests for the Aerospike online store. + +Most of the tests here are pure Python and run in any environment (they cover +the timestamp/TTL helpers, the column-oriented proto reshape, and the +write/read/admin dispatch with a mocked Aerospike client). One end-to-end test +is marked with ``@_requires_docker`` and is skipped when Docker is unavailable. +""" + +import threading +import time +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +pytest.importorskip("aerospike") + +import aerospike # noqa: E402 + +from feast import FeatureView, Field, FileSource # noqa: E402 +from feast.infra.online_stores.aerospike_online_store.aerospike import ( # noqa: E402 + AerospikeOnlineStore, + _datetime_to_epoch_ms, + _epoch_ms_to_datetime, + _resolve_ttl, +) +from feast.protos.feast.types.EntityKey_pb2 import ( + EntityKey as EntityKeyProto, # noqa: E402 +) +from feast.protos.feast.types.Value_pb2 import Value as ValueProto # noqa: E402 +from feast.repo_config import RepoConfig # noqa: E402 +from feast.types import Float64, Int64 # noqa: E402 +from feast.utils import _utc_now # noqa: E402 +from tests.universal.feature_repos.universal.online_store.aerospike import ( # noqa: E402 + AEROSPIKE_CE_IMAGE, +) +from tests.utils.cli_repo_creator import CliRunner, get_example_repo # noqa: E402 + +docker_available = False +try: + import docker + from testcontainers.core.container import DockerContainer + from testcontainers.core.waiting_utils import wait_for_logs + + try: + _docker = docker.from_env() + _docker.ping() + docker_available = True + except Exception: + pass +except ImportError: + pass + +_requires_docker = pytest.mark.skipif( + not docker_available, + reason="Docker is not available or not running. Start Docker daemon to run these tests.", +) + + +# --------------------------------------------------------------------------- +# Shared fixtures / helpers +# --------------------------------------------------------------------------- + + +def _make_fv(*field_names: str, dtype=Int64) -> FeatureView: + """Build a minimal FeatureView for conversion tests.""" + return FeatureView( + name="test_fv", + entities=[], + schema=[Field(name=n, dtype=dtype) for n in field_names], + source=FileSource(path="fake.parquet", timestamp_field="event_timestamp"), + ttl=timedelta(days=1), + ) + + +def _aerospike_repo_config(**online_store_overrides) -> RepoConfig: + base = {"type": "aerospike", "namespace": "feast"} + base.update(online_store_overrides) + return RepoConfig( + project="demo", + provider="local", + registry="/tmp/reg.db", + online_store=base, + entity_key_serialization_version=3, + ) + + +def _fake_batch_record(key: tuple, bins): + """Mimic aerospike_helpers.batch.records.BatchRecord for a successful read.""" + return SimpleNamespace( + key=key, + result=0, + record=(key, {"ttl": 0, "gen": 1}, bins) if bins is not None else None, + in_doubt=False, + ) + + +# --------------------------------------------------------------------------- +# Helpers: timestamp and TTL conversions +# --------------------------------------------------------------------------- + + +def test_datetime_helpers_round_trip_utc(): + dt = datetime(2026, 4, 20, 12, 30, 45, 123000, tzinfo=timezone.utc) + ms = _datetime_to_epoch_ms(dt) + assert _epoch_ms_to_datetime(ms) == dt + + +def test_datetime_helpers_treat_naive_as_utc(): + dt_naive = datetime(2026, 4, 20, 12, 30, 45, 123000) + dt_utc = dt_naive.replace(tzinfo=timezone.utc) + assert _datetime_to_epoch_ms(dt_naive) == _datetime_to_epoch_ms(dt_utc) + + +def test_epoch_ms_to_datetime_none_passthrough(): + assert _epoch_ms_to_datetime(None) is None + + +def test_resolve_ttl_sentinels(): + assert _resolve_ttl(None) == aerospike.TTL_NAMESPACE_DEFAULT + assert _resolve_ttl(0) == aerospike.TTL_NEVER_EXPIRE + assert _resolve_ttl(3600) == 3600 + + +def test_socket_timeout_ms_propagates_to_read_write_and_batch_policies(monkeypatch): + """``socket_timeout_ms`` must apply uniformly to read, write and batch + policies — otherwise ``max_retries`` can't fire within the total budget + (the first attempt alone consumes the whole deadline).""" + captured: dict = {} + + def fake_client(cfg): + captured["cfg"] = cfg + fake = MagicMock() + fake.connect.return_value = fake + return fake + + monkeypatch.setattr(aerospike, "client", fake_client) + + config = _aerospike_repo_config( + batch_total_timeout_ms=1500, + socket_timeout_ms=40, + read_timeout_ms=100, + write_timeout_ms=200, + ) + store = AerospikeOnlineStore() + store._get_client(config) + + policies = captured["cfg"]["policies"] + assert policies["read"]["total_timeout"] == 100 + assert policies["read"]["socket_timeout"] == 40 + assert policies["write"]["total_timeout"] == 200 + assert policies["write"]["socket_timeout"] == 40 + assert policies["batch"]["total_timeout"] == 1500 + assert policies["batch"]["socket_timeout"] == 40 + + +def test_socket_timeout_ms_omitted_when_not_set(monkeypatch): + """Unset ``socket_timeout_ms`` must leave the client's default in place, + not inject ``None`` into the policy dicts.""" + captured: dict = {} + + def fake_client(cfg): + captured["cfg"] = cfg + fake = MagicMock() + fake.connect.return_value = fake + return fake + + monkeypatch.setattr(aerospike, "client", fake_client) + + store = AerospikeOnlineStore() + store._get_client(_aerospike_repo_config()) + + policies = captured["cfg"]["policies"] + for scope in ("read", "write", "batch"): + assert "socket_timeout" not in policies[scope], ( + f"socket_timeout leaked into {scope} policy with no config override" + ) + + +# --------------------------------------------------------------------------- +# _convert_raw_docs_to_proto — same contract as MongoDB's helper +# --------------------------------------------------------------------------- + + +def test_convert_raw_docs_missing_entity(): + """Entity key absent from docs -> (None, None).""" + fv = _make_fv("score") + ts = datetime(2024, 1, 1, tzinfo=timezone.utc) + ids = [b"present", b"missing"] + docs = { + b"present": { + "features": {"test_fv": {"score": 42}}, + "event_timestamps": {"test_fv": ts}, + } + } + + results = AerospikeOnlineStore._convert_raw_docs_to_proto(ids, docs, fv) + + assert len(results) == 2 + ts_out, feats_out = results[0] + assert ts_out == ts + assert feats_out["score"].int64_val == 42 + assert results[1] == (None, None) + + +def test_convert_raw_docs_partial_doc(): + """Entity exists but one feature key is absent -> empty ValueProto for that feature.""" + fv = _make_fv("present_feat", "missing_feat") + ts = datetime(2024, 1, 1, tzinfo=timezone.utc) + ids = [b"entity1"] + docs = { + b"entity1": { + "features": {"test_fv": {"present_feat": 99}}, + "event_timestamps": {"test_fv": ts}, + } + } + + results = AerospikeOnlineStore._convert_raw_docs_to_proto(ids, docs, fv) + + assert len(results) == 1 + ts_out, feats_out = results[0] + assert ts_out == ts + assert feats_out["present_feat"].int64_val == 99 + assert feats_out["missing_feat"] == ValueProto() + + +def test_convert_raw_docs_entity_exists_but_fv_not_written(): + """Entity doc exists (written by another FV) but this FV was never written -> (None, None).""" + pricing_fv = _make_fv("price") + ts = datetime(2024, 1, 1, tzinfo=timezone.utc) + ids = [b"driver_1"] + docs = { + b"driver_1": { + "features": {"driver_stats": {"acc_rate": 0.9}}, + "event_timestamps": {"driver_stats": ts}, + } + } + + results = AerospikeOnlineStore._convert_raw_docs_to_proto(ids, docs, pricing_fv) + + assert len(results) == 1 + assert results[0] == (None, None) + + +def test_convert_raw_docs_ordering(): + """Result order matches the ids list regardless of dict insertion order in docs.""" + fv = _make_fv("score") + ts = datetime(2024, 1, 1, tzinfo=timezone.utc) + + ids = [b"entity_z", b"entity_a", b"entity_m"] + docs = { + b"entity_a": { + "features": {"test_fv": {"score": 2}}, + "event_timestamps": {"test_fv": ts}, + }, + b"entity_m": { + "features": {"test_fv": {"score": 3}}, + "event_timestamps": {"test_fv": ts}, + }, + b"entity_z": { + "features": {"test_fv": {"score": 1}}, + "event_timestamps": {"test_fv": ts}, + }, + } + + results = AerospikeOnlineStore._convert_raw_docs_to_proto(ids, docs, fv) + + assert [row[1]["score"].int64_val for row in results] == [1, 2, 3] + + +# --------------------------------------------------------------------------- +# Write path: _build_batch_writes + online_write_batch dispatch +# --------------------------------------------------------------------------- + + +def _entity_key(join_key: str, value: int) -> EntityKeyProto: + return EntityKeyProto( + join_keys=[join_key], entity_values=[ValueProto(int64_val=value)] + ) + + +def test_build_batch_writes_produces_three_ops_with_created_ts(): + config = _aerospike_repo_config(ttl_seconds=3600) + fv = SimpleNamespace(name="driver_stats") + ts = datetime(2026, 4, 20, 12, 30, 45, tzinfo=timezone.utc) + row = ( + _entity_key("driver_id", 1), + { + "rating": ValueProto(double_val=4.91), + "trips_last_7d": ValueProto(int64_val=132), + }, + ts, + ts, + ) + + batch = AerospikeOnlineStore._build_batch_writes( + config, fv, [row], namespace="feast", set_name="demo_latest" + ) + + assert len(batch.batch_records) == 1 + bw = batch.batch_records[0] + assert bw.key[:2] == ("feast", "demo_latest") + # Must be bytearray, not bytes: the Aerospike Python C client rejects + # bytes user keys ("Key is invalid") and silently hashes only the first + # byte inside batch_operate/batch_read — causing digest collisions. + assert isinstance(bw.key[2], bytearray) + assert bw.meta == {"ttl": 3600} + # No explicit per-record policy: writes rely on the client-level default + # (POLICY_KEY_DIGEST), which doesn't persist the serialized entity key + # server-side — batch_operate preserves request order so the stored key + # has no functional use on the read path. + assert bw.policy is None + assert len(bw.ops) == 3 + bin_names = [op["bin"] for op in bw.ops] + assert bin_names == ["features", "event_ts", "created_ts"] + # Map CDTs must be created with MAP_KEY_ORDERED so key lookups on reads + # and the update() background scan stay O(log N). + for op in bw.ops: + if op["bin"] in ("features", "event_ts"): + assert op["map_policy"] == {"map_order": aerospike.MAP_KEY_ORDERED} + + +def test_build_batch_writes_omits_created_ts_when_none(): + config = _aerospike_repo_config() + fv = SimpleNamespace(name="driver_stats") + row = ( + _entity_key("driver_id", 1), + {"rating": ValueProto(double_val=4.91)}, + datetime(2026, 4, 20, tzinfo=timezone.utc), + None, + ) + + batch = AerospikeOnlineStore._build_batch_writes( + config, fv, [row], namespace="feast", set_name="demo_latest" + ) + + ops = batch.batch_records[0].ops + assert len(ops) == 2 + assert {op["bin"] for op in ops} == {"features", "event_ts"} + + +def test_build_batch_writes_ttl_sentinels(): + config = _aerospike_repo_config(ttl_seconds=0) + fv = SimpleNamespace(name="fv") + row = ( + _entity_key("id", 1), + {"x": ValueProto(int64_val=1)}, + datetime(2026, 1, 1, tzinfo=timezone.utc), + None, + ) + + batch = AerospikeOnlineStore._build_batch_writes( + config, fv, [row], namespace="feast", set_name="set" + ) + assert batch.batch_records[0].meta == {"ttl": aerospike.TTL_NEVER_EXPIRE} + + +def test_online_write_batch_dispatches_to_client(): + config = _aerospike_repo_config() + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + row = ( + _entity_key("id", 1), + {"x": ValueProto(int64_val=1)}, + datetime(2026, 1, 1, tzinfo=timezone.utc), + None, + ) + progress_calls: list[int] = [] + store.online_write_batch(config, fv, [row], progress=progress_calls.append) + + assert fake_client.batch_write.called + assert progress_calls == [1] + + +def test_online_write_batch_empty_short_circuits(): + config = _aerospike_repo_config() + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + progress_calls: list[int] = [] + store.online_write_batch(config, fv, [], progress=progress_calls.append) + + assert not fake_client.batch_write.called + assert progress_calls == [0] + + +def test_chunked_splits_without_partial_tail(): + chunks = list(AerospikeOnlineStore._chunked([1, 2, 3, 4, 5], 2)) + assert chunks == [[1, 2], [3, 4], [5]] + + +def test_get_client_thread_safe_single_connect(monkeypatch): + """Concurrent first callers must share one client, not leak connections.""" + connect_count = 0 + connect_lock = threading.Lock() + + def fake_client(cfg): + fake = MagicMock() + + def connect(): + nonlocal connect_count + with connect_lock: + connect_count += 1 + time.sleep(0.05) + return fake + + fake.connect = connect + return fake + + monkeypatch.setattr(aerospike, "client", fake_client) + store = AerospikeOnlineStore() + config = _aerospike_repo_config() + barrier = threading.Barrier(8) + + def worker(): + barrier.wait() + store._get_client(config) + + threads = [threading.Thread(target=worker) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + assert connect_count == 1 + + +def test_online_write_batch_chunks_at_batch_max_records(): + config = _aerospike_repo_config(batch_max_records=100) + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + ts = datetime(2026, 1, 1, tzinfo=timezone.utc) + rows = [ + ( + _entity_key("id", i), + {"x": ValueProto(int64_val=i)}, + ts, + None, + ) + for i in range(250) + ] + progress_calls: list[int] = [] + store.online_write_batch(config, fv, rows, progress=progress_calls.append) + + assert fake_client.batch_write.call_count == 3 + sizes = [ + len(call.args[0].batch_records) + for call in fake_client.batch_write.call_args_list + ] + assert sizes == [100, 100, 50] + assert progress_calls == [100, 200, 250] + + +# --------------------------------------------------------------------------- +# Read path: online_read dispatches and converts via batch_operate +# --------------------------------------------------------------------------- + + +def _read_feature_view() -> SimpleNamespace: + """Minimal FV object exposing .name and .features with dtype mappings.""" + return SimpleNamespace( + name="driver_stats", + features=[ + SimpleNamespace(name="rating", dtype=Float64), + SimpleNamespace(name="trips_last_7d", dtype=Int64), + ], + ) + + +def test_online_read_happy_path_with_projection_and_ordering(): + config = _aerospike_repo_config() + fv = _read_feature_view() + store = AerospikeOnlineStore() + ts = datetime(2026, 4, 20, 12, 30, 45, 123000, tzinfo=timezone.utc) + + ek1 = _entity_key("driver_id", 1) + ek2 = _entity_key("driver_id", 2) + ek3 = _entity_key("driver_id", 3) + key1 = store._aerospike_key(config, ek1) + key2 = store._aerospike_key(config, ek2) + key3 = store._aerospike_key(config, ek3) + + def fake_batch_operate(keys, ops): + assert keys == [key1, key2, key3] + assert len(ops) == 2 + assert ops[0]["bin"] == "features" + assert ops[1]["bin"] == "event_ts" + # Aerospike's batch_operate preserves input order; the middle record + # is simulated as missing to verify we still emit (None, None) in + # the correct slot. + br1 = _fake_batch_record( + key1, + { + "features": {"rating": 4.91, "trips_last_7d": 132}, + "event_ts": _datetime_to_epoch_ms(ts), + }, + ) + br2 = _fake_batch_record(key2, None) # missing record + br3 = _fake_batch_record( + key3, + { + "features": {"rating": 3.75, "trips_last_7d": 42}, + "event_ts": _datetime_to_epoch_ms(ts), + }, + ) + return SimpleNamespace(batch_records=[br1, br2, br3]) + + fake_client = MagicMock() + fake_client.batch_operate.side_effect = fake_batch_operate + store._client = fake_client + + results = store.online_read(config, fv, [ek1, ek2, ek3]) + + assert len(results) == 3 + ts0, feats0 = results[0] + assert ts0 == ts + assert abs(feats0["rating"].double_val - 4.91) < 1e-9 + assert feats0["trips_last_7d"].int64_val == 132 + assert results[1] == (None, None) + ts2, feats2 = results[2] + assert ts2 == ts + assert abs(feats2["rating"].double_val - 3.75) < 1e-9 + + +def test_online_read_empty_keys_returns_empty(): + store = AerospikeOnlineStore() + store._client = MagicMock() + fv = _read_feature_view() + assert store.online_read(_aerospike_repo_config(), fv, []) == [] + assert not store._client.batch_operate.called + + +def test_online_read_chunks_at_batch_max_records(): + config = _aerospike_repo_config(batch_max_records=100) + fv = _read_feature_view() + store = AerospikeOnlineStore() + entity_keys = [_entity_key("driver_id", i) for i in range(250)] + ts = datetime(2026, 1, 1, tzinfo=timezone.utc) + + def fake_batch_operate(keys, ops): + return SimpleNamespace( + batch_records=[ + _fake_batch_record( + key, + { + "features": {"rating": 4.0, "trips_last_7d": 10}, + "event_ts": _datetime_to_epoch_ms(ts), + }, + ) + for key in keys + ] + ) + + fake_client = MagicMock() + fake_client.batch_operate.side_effect = fake_batch_operate + store._client = fake_client + + results = store.online_read(config, fv, entity_keys) + + assert fake_client.batch_operate.call_count == 3 + sizes = [len(call.args[0]) for call in fake_client.batch_operate.call_args_list] + assert sizes == [100, 100, 50] + assert len(results) == 250 + assert all(feats is not None for _, feats in results) + + +def test_online_read_record_exists_but_fv_not_present_returns_none(): + config = _aerospike_repo_config() + fv = _read_feature_view() + store = AerospikeOnlineStore() + ek = _entity_key("driver_id", 1) + key = store._aerospike_key(config, ek) + + def fake_batch_operate(keys, ops): + return SimpleNamespace( + batch_records=[ + _fake_batch_record(key, {"features": None, "event_ts": None}) + ] + ) + + fake_client = MagicMock() + fake_client.batch_operate.side_effect = fake_batch_operate + store._client = fake_client + + results = store.online_read(config, fv, [ek]) + assert results == [(None, None)] + + +def _fake_error_record(key: tuple, result_code: int): + """Batch response for an error case — no ``record`` but a non-OK result.""" + return SimpleNamespace(key=key, result=result_code, record=None, in_doubt=False) + + +def test_online_read_record_not_found_returns_none(): + """``br.result == 2`` (RECORD_NOT_FOUND) is a genuine miss, not an error.""" + config = _aerospike_repo_config() + fv = _read_feature_view() + store = AerospikeOnlineStore() + ek = _entity_key("driver_id", 1) + key = store._aerospike_key(config, ek) + + fake_client = MagicMock() + fake_client.batch_operate.return_value = SimpleNamespace( + batch_records=[_fake_error_record(key, result_code=2)] + ) + store._client = fake_client + + assert store.online_read(config, fv, [ek]) == [(None, None)] + + +def test_online_read_op_not_applicable_returns_none(): + """``br.result == 26`` (OP_NOT_APPLICABLE) happens when the nested FV + submap is absent — also a miss under the OnlineStore contract, not an + error.""" + config = _aerospike_repo_config() + fv = _read_feature_view() + store = AerospikeOnlineStore() + ek = _entity_key("driver_id", 1) + key = store._aerospike_key(config, ek) + + fake_client = MagicMock() + fake_client.batch_operate.return_value = SimpleNamespace( + batch_records=[_fake_error_record(key, result_code=26)] + ) + store._client = fake_client + + assert store.online_read(config, fv, [ek]) == [(None, None)] + + +def test_online_read_raises_on_transient_error(): + """Any other non-OK result (e.g. 9 = TIMEOUT) must be surfaced as an error. + + A silent null-return on a transient timeout is the root-cause class of + bug that shows up weeks later as a model-quality regression; pinning + this with a test keeps the read path loud on failure. + """ + config = _aerospike_repo_config() + fv = _read_feature_view() + store = AerospikeOnlineStore() + ek = _entity_key("driver_id", 1) + key = store._aerospike_key(config, ek) + + fake_client = MagicMock() + fake_client.batch_operate.return_value = SimpleNamespace( + batch_records=[_fake_error_record(key, result_code=9)] # TIMEOUT + ) + store._client = fake_client + + with pytest.raises(RuntimeError, match="non-OK status"): + store.online_read(config, fv, [ek]) + + +def test_online_read_projects_requested_features_server_side(): + """``requested_features`` must drive a ``map_get_by_key_list`` projection + nested into the feature-view submap via ``cdt_ctx``; the unordered + ``MAP_RETURN_KEY_VALUE`` response is a flat ``[k1,v1,k2,v2]`` list.""" + config = _aerospike_repo_config() + fv = _read_feature_view() + store = AerospikeOnlineStore() + ek = _entity_key("driver_id", 1) + key = store._aerospike_key(config, ek) + ts = datetime(2026, 4, 20, 12, 30, 45, tzinfo=timezone.utc) + + captured_ops: list = [] + + def fake_batch_operate(keys, ops): + captured_ops.extend(ops) + # Projected features come back as a flat [k,v,k,v] list rather than + # a dict — the store must normalize this before feeding the row + # reshape helper. + return SimpleNamespace( + batch_records=[ + _fake_batch_record( + key, + { + "features": ["rating", 4.91], + "event_ts": _datetime_to_epoch_ms(ts), + }, + ) + ] + ) + + fake_client = MagicMock() + fake_client.batch_operate.side_effect = fake_batch_operate + store._client = fake_client + + results = store.online_read(config, fv, [ek], requested_features=["rating"]) + + assert len(results) == 1 + ts_out, feats = results[0] + assert ts_out == ts + assert abs(feats["rating"].double_val - 4.91) < 1e-9 + assert feats["trips_last_7d"] == ValueProto() + features_op = captured_ops[0] + assert features_op["op"] == aerospike.OP_MAP_GET_BY_KEY_LIST + assert features_op["bin"] == "features" + assert features_op["val"] == ["rating"] + assert features_op["return_type"] == aerospike.MAP_RETURN_KEY_VALUE + # The ctx must nest the projection inside the FV's submap so the server + # only ships the requested columns over the wire. + assert features_op.get("ctx"), "projection op is missing ctx" + + +def test_normalize_projected_features_handles_all_payload_shapes(): + """The shape of the ``features`` payload depends on which op produced it; + the helper must accept all of them.""" + assert AerospikeOnlineStore._normalize_projected_features(None) is None + assert AerospikeOnlineStore._normalize_projected_features([]) == {} + assert AerospikeOnlineStore._normalize_projected_features(["a", 1, "b", 2]) == { + "a": 1, + "b": 2, + } + assert AerospikeOnlineStore._normalize_projected_features({"a": 1, "b": 2}) == { + "a": 1, + "b": 2, + } + + +def test_online_write_batch_raises_on_per_record_error(): + """A partial batch failure (one record's result != 0) must raise rather + than silently succeed — otherwise downstream sees "model saw stale + features" weeks after the fact.""" + config = _aerospike_repo_config() + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + + def fake_batch_write(batch): + # Simulate a one-partition timeout on the single record. + batch.batch_records[0].result = 9 # TIMEOUT + + fake_client.batch_write.side_effect = fake_batch_write + store._client = fake_client + + row = ( + _entity_key("id", 1), + {"x": ValueProto(int64_val=1)}, + datetime(2026, 1, 1, tzinfo=timezone.utc), + None, + ) + with pytest.raises(RuntimeError, match="non-OK"): + store.online_write_batch(config, fv, [row], progress=None) + + +# --------------------------------------------------------------------------- +# Admin paths: update / teardown +# --------------------------------------------------------------------------- + + +def test_update_no_op_when_nothing_to_delete(): + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + store.update(_aerospike_repo_config(), [], [], [], [], partial=False) + assert not fake_client.scan.called + + +def test_update_single_fv_issues_single_background_scan(): + store = AerospikeOnlineStore() + fake_client = MagicMock() + fake_scan = MagicMock() + fake_client.scan.return_value = fake_scan + store._client = fake_client + + store.update( + _aerospike_repo_config(), + [SimpleNamespace(name="old_fv")], + [], + [], + [], + partial=False, + ) + + assert fake_client.scan.call_args[0] == ("feast", "demo_latest") + ops = fake_scan.add_ops.call_args[0][0] + assert len(ops) == 2 + assert {op["bin"] for op in ops} == {"features", "event_ts"} + for op in ops: + assert op["key"] == "old_fv" + assert op["return_type"] == aerospike.MAP_RETURN_NONE + fake_scan.execute_background.assert_called_once() + + +def test_update_multi_fv_coalesces_into_one_scan(): + store = AerospikeOnlineStore() + fake_client = MagicMock() + fake_scan = MagicMock() + fake_client.scan.return_value = fake_scan + store._client = fake_client + + tables = [SimpleNamespace(name=n) for n in ("a", "b", "c")] + store.update(_aerospike_repo_config(), tables, [], [], [], partial=False) + + fake_client.scan.assert_called_once() + ops = fake_scan.add_ops.call_args[0][0] + assert len(ops) == 6 + assert {op["bin"] for op in ops} == {"features", "event_ts"} + assert {op["key"] for op in ops} == {"a", "b", "c"} + fake_scan.execute_background.assert_called_once() + + +def test_update_rejects_non_aerospike_config(): + wrong_config = RepoConfig( + project="demo", + provider="local", + registry="/tmp/reg.db", + online_store={"type": "sqlite", "path": "/tmp/online.db"}, + entity_key_serialization_version=3, + ) + store = AerospikeOnlineStore() + with pytest.raises(RuntimeError): + store.update( + wrong_config, + [SimpleNamespace(name="x")], + [], + [], + [], + partial=False, + ) + + +def test_teardown_truncates_and_closes_client(): + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + store.teardown(_aerospike_repo_config(), [], []) + assert fake_client.truncate.call_args[0] == ("feast", "demo_latest", 0) + fake_client.close.assert_called_once() + assert store._client is None + + +def test_teardown_rejects_non_aerospike_config(): + wrong_config = RepoConfig( + project="demo", + provider="local", + registry="/tmp/reg.db", + online_store={"type": "sqlite", "path": "/tmp/online.db"}, + entity_key_serialization_version=3, + ) + store = AerospikeOnlineStore() + with pytest.raises(RuntimeError): + store.teardown(wrong_config, [], []) + + +# --------------------------------------------------------------------------- +# Per-feature-view namespace + set overrides +# --------------------------------------------------------------------------- +# These let a deployment pin individual feature views to RAM-only / SSD-backed +# namespaces or isolate one view in its own set without splitting projects. +# Coverage: +# * the ``_namespace_for_fv`` / ``_set_name`` resolvers +# * read + write paths actually use the resolved ns/set on the wire +# * ``update()`` issues one scan per (ns, set) group +# * ``teardown()`` truncates every unique (ns, set) pair + + +def test_set_name_default_when_no_overrides(): + config = _aerospike_repo_config() + store = AerospikeOnlineStore() + assert store._set_name(config) == "demo_latest" + assert store._set_name(config, "any_fv") == "demo_latest" + + +def test_set_name_resolves_per_fv_override(): + config = _aerospike_repo_config( + set_overrides={"hot_fv": "demo_hot", "cold_fv": "demo_cold"} + ) + store = AerospikeOnlineStore() + assert store._set_name(config, "hot_fv") == "demo_hot" + assert store._set_name(config, "cold_fv") == "demo_cold" + # Unmapped FV still falls back to the project-level default. + assert store._set_name(config, "other_fv") == "demo_latest" + # Calling without an FV name explicitly asks for the project default. + assert store._set_name(config) == "demo_latest" + + +def test_namespace_for_fv_default_when_no_overrides(): + config = _aerospike_repo_config() + store = AerospikeOnlineStore() + assert store._namespace_for_fv(config) == "feast" + assert store._namespace_for_fv(config, "any_fv") == "feast" + + +def test_namespace_for_fv_resolves_per_fv_override(): + config = _aerospike_repo_config( + namespace_overrides={"hot_fv": "feast_ram", "cold_fv": "feast_ssd"} + ) + store = AerospikeOnlineStore() + assert store._namespace_for_fv(config, "hot_fv") == "feast_ram" + assert store._namespace_for_fv(config, "cold_fv") == "feast_ssd" + assert store._namespace_for_fv(config, "other_fv") == "feast" + + +def test_aerospike_key_honours_overrides_when_fv_passed(): + """``_aerospike_key`` is the choke point for everything that builds an + Aerospike (ns, set, user_key) tuple — it must respect the overrides + when an FV name is provided, and stay backwards-compatible without.""" + config = _aerospike_repo_config( + namespace_overrides={"hot_fv": "feast_ram"}, + set_overrides={"hot_fv": "demo_hot"}, + ) + store = AerospikeOnlineStore() + ek = _entity_key("driver_id", 1) + + default_ns, default_set, _ = store._aerospike_key(config, ek) + assert (default_ns, default_set) == ("feast", "demo_latest") + + hot_ns, hot_set, _ = store._aerospike_key(config, ek, fv_name="hot_fv") + assert (hot_ns, hot_set) == ("feast_ram", "demo_hot") + + +def test_online_write_batch_writes_to_per_fv_ns_and_set(): + """The wire keys produced by online_write_batch must use the FV's + resolved ns + set, not the store-level defaults.""" + config = _aerospike_repo_config( + namespace_overrides={"driver_stats": "feast_ram"}, + set_overrides={"driver_stats": "demo_hot"}, + ) + fv = SimpleNamespace(name="driver_stats") + store = AerospikeOnlineStore() + fake_client = MagicMock() + captured: dict = {} + + def fake_batch_write(batch): + captured["batch"] = batch + + fake_client.batch_write.side_effect = fake_batch_write + store._client = fake_client + + row = ( + _entity_key("id", 1), + {"x": ValueProto(int64_val=1)}, + datetime(2026, 1, 1, tzinfo=timezone.utc), + None, + ) + store.online_write_batch(config, fv, [row], progress=None) + + bw = captured["batch"].batch_records[0] + assert bw.key[:2] == ("feast_ram", "demo_hot") + + +def test_online_read_uses_per_fv_ns_and_set(): + """The ``batch_operate`` keys must use the FV's resolved ns + set.""" + config = _aerospike_repo_config( + namespace_overrides={"driver_stats": "feast_ram"}, + set_overrides={"driver_stats": "demo_hot"}, + ) + fv = _read_feature_view() # name="driver_stats" + store = AerospikeOnlineStore() + captured: dict = {} + + def fake_batch_operate(keys, ops): + captured["keys"] = keys + return SimpleNamespace(batch_records=[_fake_batch_record(keys[0], None)]) + + fake_client = MagicMock() + fake_client.batch_operate.side_effect = fake_batch_operate + store._client = fake_client + + store.online_read(config, fv, [_entity_key("driver_id", 1)]) + assert len(captured["keys"]) == 1 + assert captured["keys"][0][:2] == ("feast_ram", "demo_hot") + + +def test_update_groups_dropped_fvs_by_resolved_ns_and_set(): + """Dropped feature views in different (ns, set) buckets must produce + one background scan per bucket — a single combined scan would either + miss records or pointlessly scan unrelated namespaces.""" + config = _aerospike_repo_config( + namespace_overrides={"a": "ns_x"}, # b, c land on default ns + set_overrides={"b": "set_y"}, # a, c land on default set + ) + store = AerospikeOnlineStore() + fake_client = MagicMock() + + fake_scans: list = [] + + def make_scan(ns, set_name): + s = MagicMock(name=f"scan_{ns}_{set_name}") + s._ns = ns + s._set = set_name + fake_scans.append(s) + return s + + fake_client.scan.side_effect = make_scan + store._client = fake_client + + tables = [ + SimpleNamespace(name="a"), # (ns_x, demo_latest) + SimpleNamespace(name="b"), # (feast, set_y) + SimpleNamespace(name="c"), # (feast, demo_latest) + ] + store.update(config, tables, [], [], [], partial=False) + + targets = {(s._ns, s._set) for s in fake_scans} + assert targets == { + ("ns_x", "demo_latest"), + ("feast", "set_y"), + ("feast", "demo_latest"), + } + # Every group's scan must execute and ship the matching FVs' remove ops. + by_target = {(s._ns, s._set): s for s in fake_scans} + for tgt, fv_name in [ + (("ns_x", "demo_latest"), "a"), + (("feast", "set_y"), "b"), + (("feast", "demo_latest"), "c"), + ]: + scan = by_target[tgt] + ops = scan.add_ops.call_args[0][0] + assert {op["bin"] for op in ops} == {"features", "event_ts"} + assert all(op["key"] == fv_name for op in ops) + scan.execute_background.assert_called_once() + + +def test_update_coalesces_fvs_sharing_a_resolved_target(): + """Two dropped FVs that resolve to the same (ns, set) must share one + scan — the per-FV grouping is "by target", not "by FV name".""" + config = _aerospike_repo_config( + namespace_overrides={"a": "ns_x", "b": "ns_x"}, + set_overrides={"a": "set_y", "b": "set_y"}, + ) + store = AerospikeOnlineStore() + fake_client = MagicMock() + fake_scan = MagicMock() + fake_client.scan.return_value = fake_scan + store._client = fake_client + + store.update( + config, + [SimpleNamespace(name="a"), SimpleNamespace(name="b")], + [], + [], + [], + partial=False, + ) + + fake_client.scan.assert_called_once_with("ns_x", "set_y") + ops = fake_scan.add_ops.call_args[0][0] + # 4 ops total: features+event_ts for each of the 2 FVs. + assert len(ops) == 4 + assert {op["key"] for op in ops} == {"a", "b"} + + +def test_teardown_truncates_default_plus_every_overridden_pair(): + config = _aerospike_repo_config( + namespace_overrides={"hot": "feast_ram"}, + set_overrides={"isolated": "demo_iso"}, + ) + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + tables = [ + SimpleNamespace(name="hot"), # (feast_ram, demo_latest) + SimpleNamespace(name="isolated"), # (feast, demo_iso) + SimpleNamespace(name="default"), # (feast, demo_latest) + ] + store.teardown(config, tables, []) + + truncated = {tuple(call.args[:2]) for call in fake_client.truncate.call_args_list} + assert truncated == { + ("feast", "demo_latest"), + ("feast_ram", "demo_latest"), + ("feast", "demo_iso"), + } + # cutoff is always 0 ("drop everything regardless of last-update time") + for call in fake_client.truncate.call_args_list: + assert call.args[2] == 0 + fake_client.close.assert_called_once() + + +def test_teardown_with_empty_tables_still_clears_default_pair(): + """A teardown call with an empty tables list (e.g. a bare ``feast + teardown`` on a repo whose registry is empty) must still clear the + store-default location, otherwise stale data lingers.""" + config = _aerospike_repo_config( + namespace_overrides={"unused": "feast_ram"}, + ) + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + store.teardown(config, [], []) + + truncated = {tuple(call.args[:2]) for call in fake_client.truncate.call_args_list} + assert truncated == {("feast", "demo_latest")} + + +# --------------------------------------------------------------------------- +# Prewriting hook +# --------------------------------------------------------------------------- +# Hook-target functions must be module-level so the import-string resolver +# can find them via importlib + getattr. They're referenced by f"{__name__}.X". + + +def _hook_uppercase_string_features(config, table, data): # noqa: ARG001 + """Sample hook: uppercase every string ValueProto.string_val on every row.""" + new_data = [] + for entity_key, values, ts, created in data: + new_values = {} + for k, v in values.items(): + if v.HasField("string_val"): + new_values[k] = ValueProto(string_val=v.string_val.upper()) + else: + new_values[k] = v + new_data.append((entity_key, new_values, ts, created)) + return new_data + + +def _hook_drop_all_rows(config, table, data): # noqa: ARG001 + """Sample hook that filters every row — exercises the post-hook empty-batch path.""" + return [] + + +def _hook_raises(config, table, data): # noqa: ARG001 + raise ValueError("hook intentionally failed") + + +_NOT_A_FUNCTION = 42 + + +def test_prewriting_hook_unset_returns_none(): + config = _aerospike_repo_config() # no prewriting_hook + store = AerospikeOnlineStore() + assert store._resolve_prewriting_hook(config) is None + # Must clear any previously-cached value so a config swap takes effect. + store._prewriting_hook = lambda *a, **kw: None # type: ignore[assignment] + store._prewriting_hook_spec = "stale.spec" + store._resolve_prewriting_hook(config) + assert store._prewriting_hook is None + assert store._prewriting_hook_spec is None + + +def test_prewriting_hook_resolves_and_caches(): + spec = f"{__name__}._hook_uppercase_string_features" + config = _aerospike_repo_config(prewriting_hook=spec) + store = AerospikeOnlineStore() + + hook = store._resolve_prewriting_hook(config) + assert hook is _hook_uppercase_string_features + + # Second call must hit the cache (no re-import). + with patch( + "feast.infra.online_stores.aerospike_online_store.aerospike.importlib" + ) as m: + again = store._resolve_prewriting_hook(config) + assert again is hook + assert not m.import_module.called + + +def test_prewriting_hook_recompiles_after_spec_change(): + """If the config is rebound to a new hook string, the resolver must + re-resolve instead of returning the cached old callable.""" + store = AerospikeOnlineStore() + + cfg1 = _aerospike_repo_config( + prewriting_hook=f"{__name__}._hook_uppercase_string_features" + ) + assert store._resolve_prewriting_hook(cfg1) is _hook_uppercase_string_features + + cfg2 = _aerospike_repo_config(prewriting_hook=f"{__name__}._hook_drop_all_rows") + assert store._resolve_prewriting_hook(cfg2) is _hook_drop_all_rows + + +def test_prewriting_hook_bad_format_raises(): + config = _aerospike_repo_config(prewriting_hook="no_dot_here") + store = AerospikeOnlineStore() + with pytest.raises(ValueError, match="fully qualified import path"): + store._resolve_prewriting_hook(config) + + +def test_prewriting_hook_missing_module_raises(): + config = _aerospike_repo_config(prewriting_hook="definitely_not_a_real.module.fn") + store = AerospikeOnlineStore() + with pytest.raises(ValueError, match="could not import module"): + store._resolve_prewriting_hook(config) + + +def test_prewriting_hook_missing_attribute_raises(): + config = _aerospike_repo_config(prewriting_hook=f"{__name__}._does_not_exist") + store = AerospikeOnlineStore() + with pytest.raises(ValueError, match="has no attribute"): + store._resolve_prewriting_hook(config) + + +def test_prewriting_hook_non_callable_raises(): + config = _aerospike_repo_config(prewriting_hook=f"{__name__}._NOT_A_FUNCTION") + store = AerospikeOnlineStore() + with pytest.raises(TypeError, match="non-callable"): + store._resolve_prewriting_hook(config) + + +def test_prewriting_hook_transforms_data_before_write(): + """End-to-end: the hook's output is what lands on the wire, not the + caller-supplied data.""" + config = _aerospike_repo_config( + prewriting_hook=f"{__name__}._hook_uppercase_string_features" + ) + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + captured: dict = {} + + def fake_batch_write(batch): + captured["batch"] = batch + + fake_client.batch_write.side_effect = fake_batch_write + store._client = fake_client + + row = ( + _entity_key("id", 1), + {"name": ValueProto(string_val="alice")}, + datetime(2026, 1, 1, tzinfo=timezone.utc), + None, + ) + store.online_write_batch(config, fv, [row], progress=None) + + bw = captured["batch"].batch_records[0] + map_put_op = next(op for op in bw.ops if op["bin"] == "features") + # The hook upper-cases string features. ``map_put_items`` stores its + # payload under the ``val`` key (mapping FV name -> feature submap). + fv_map = map_put_op["val"]["fv"] + assert fv_map == {"name": "ALICE"} + + +def test_prewriting_hook_returning_empty_short_circuits(): + """A hook that filters every row must skip the wire call and report + progress=0 — same shape as the empty-input fast path.""" + config = _aerospike_repo_config(prewriting_hook=f"{__name__}._hook_drop_all_rows") + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + progress_calls: list[int] = [] + row = ( + _entity_key("id", 1), + {"x": ValueProto(int64_val=1)}, + datetime(2026, 1, 1, tzinfo=timezone.utc), + None, + ) + store.online_write_batch(config, fv, [row], progress=progress_calls.append) + + assert not fake_client.batch_write.called + assert progress_calls == [0] + + +def test_prewriting_hook_propagates_exceptions(): + """A raising hook must fail the whole batch — there's no per-row + fallback. ``batch_write`` must never be called.""" + config = _aerospike_repo_config(prewriting_hook=f"{__name__}._hook_raises") + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + row = ( + _entity_key("id", 1), + {"x": ValueProto(int64_val=1)}, + datetime(2026, 1, 1, tzinfo=timezone.utc), + None, + ) + with pytest.raises(ValueError, match="hook intentionally failed"): + store.online_write_batch(config, fv, [row], progress=None) + assert not fake_client.batch_write.called + + +def test_prewriting_hook_skipped_for_empty_input(): + """The hook resolver must not even be invoked when ``data`` is empty — + avoids paying the import cost for no-op writes.""" + config = _aerospike_repo_config( + prewriting_hook="should.never.resolve" # would raise if resolved + ) + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + progress_calls: list[int] = [] + store.online_write_batch(config, fv, [], progress=progress_calls.append) + assert not fake_client.batch_write.called + assert progress_calls == [0] + + +# --------------------------------------------------------------------------- +# Async wrappers (run_in_executor) +# --------------------------------------------------------------------------- + + +async def test_online_write_batch_async_delegates_to_sync(): + """The async write path must produce identical side-effects to the sync one.""" + config = _aerospike_repo_config() + fv = SimpleNamespace(name="fv") + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + row = ( + _entity_key("id", 1), + {"x": ValueProto(int64_val=1)}, + datetime(2026, 1, 1, tzinfo=timezone.utc), + None, + ) + progress_calls: list[int] = [] + await store.online_write_batch_async( + config, fv, [row], progress=progress_calls.append + ) + + assert fake_client.batch_write.called + assert progress_calls == [1] + + +async def test_online_read_async_returns_same_shape_as_sync(): + config = _aerospike_repo_config() + fv = _read_feature_view() + store = AerospikeOnlineStore() + ts = datetime(2026, 4, 20, tzinfo=timezone.utc) + ek = _entity_key("driver_id", 1) + key = store._aerospike_key(config, ek) + + def fake_batch_operate(keys, ops): + return SimpleNamespace( + batch_records=[ + _fake_batch_record( + key, + { + "features": {"rating": 4.91, "trips_last_7d": 132}, + "event_ts": _datetime_to_epoch_ms(ts), + }, + ) + ] + ) + + fake_client = MagicMock() + fake_client.batch_operate.side_effect = fake_batch_operate + store._client = fake_client + + results = await store.online_read_async(config, fv, [ek]) + assert len(results) == 1 + ts_out, feats = results[0] + assert ts_out == ts + assert feats["trips_last_7d"].int64_val == 132 + + +async def test_initialize_pre_warms_client(): + """initialize() must cause the client to connect without needing a read/write.""" + config = _aerospike_repo_config() + store = AerospikeOnlineStore() + assert store._client is None + + sentinel = MagicMock(name="warm_client") + store._get_client = MagicMock(return_value=sentinel) # type: ignore[assignment] + + await store.initialize(config) + store._get_client.assert_called_once_with(config) + + +async def test_close_is_noop_without_client(): + store = AerospikeOnlineStore() + assert store._client is None + await store.close() # must not raise + + +async def test_close_releases_client(): + store = AerospikeOnlineStore() + fake_client = MagicMock() + store._client = fake_client + + await store.close() + + fake_client.close.assert_called_once() + assert store._client is None + + +# --------------------------------------------------------------------------- +# End-to-end integration test — requires Docker +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def aerospike_container(): + """Start a real Aerospike CE container for end-to-end testing.""" + container = DockerContainer(AEROSPIKE_CE_IMAGE).with_exposed_ports("3000") + container.start() + wait_for_logs(container=container, predicate="migrations: complete", timeout=60) + yield container + container.stop() + + +@pytest.fixture +def aerospike_online_store_config(aerospike_container): + port = int(aerospike_container.get_exposed_port("3000")) + return { + "type": "aerospike", + "hosts": [("127.0.0.1", port)], + "namespace": "test", # default namespace shipped in the CE image + } + + +@_requires_docker +def test_aerospike_online_features(aerospike_online_store_config): + """Full round-trip: write via the provider, read via the feature store API.""" + runner = CliRunner() + with runner.local_repo( + get_example_repo("example_feature_repo_1.py"), + offline_store="file", + online_store="aerospike", + teardown=False, # container torn down by fixture + ) as store: + # Patch in the live container's port. + store.config.online_store.hosts = aerospike_online_store_config["hosts"] + store.config.online_store.namespace = aerospike_online_store_config["namespace"] + + driver_locations_fv = store.get_feature_view(name="driver_locations") + customer_profile_fv = store.get_feature_view(name="customer_profile") + customer_driver_combined_fv = store.get_feature_view( + name="customer_driver_combined" + ) + + provider = store._get_provider() + + driver_key = EntityKeyProto( + join_keys=["driver_id"], entity_values=[ValueProto(int64_val=1)] + ) + provider.online_write_batch( + config=store.config, + table=driver_locations_fv, + data=[ + ( + driver_key, + { + "lat": ValueProto(double_val=0.1), + "lon": ValueProto(string_val="1.0"), + }, + _utc_now(), + _utc_now(), + ) + ], + progress=None, + ) + + customer_key = EntityKeyProto( + join_keys=["customer_id"], entity_values=[ValueProto(string_val="5")] + ) + provider.online_write_batch( + config=store.config, + table=customer_profile_fv, + data=[ + ( + customer_key, + { + "avg_orders_day": ValueProto(float_val=1.0), + "name": ValueProto(string_val="John"), + "age": ValueProto(int64_val=3), + }, + _utc_now(), + _utc_now(), + ) + ], + progress=None, + ) + + combined_key = EntityKeyProto( + join_keys=["customer_id", "driver_id"], + entity_values=[ValueProto(string_val="5"), ValueProto(int64_val=1)], + ) + provider.online_write_batch( + config=store.config, + table=customer_driver_combined_fv, + data=[ + ( + combined_key, + {"trips": ValueProto(int64_val=7)}, + _utc_now(), + _utc_now(), + ) + ], + progress=None, + ) + + result = store.get_online_features( + features=[ + "driver_locations:lon", + "customer_profile:avg_orders_day", + "customer_profile:name", + "customer_driver_combined:trips", + ], + entity_rows=[ + {"driver_id": 1, "customer_id": "5"}, + {"driver_id": 1, "customer_id": 5}, + ], + full_feature_names=False, + ).to_dict() + + assert result["driver_id"] == [1, 1] + assert result["customer_id"] == ["5", "5"] + assert result["lon"] == ["1.0", "1.0"] + assert result["avg_orders_day"] == [1.0, 1.0] + assert result["name"] == ["John", "John"] + assert result["trips"] == [7, 7] + + missing = store.get_online_features( + features=["customer_driver_combined:trips"], + entity_rows=[{"driver_id": 0, "customer_id": 0}], + full_feature_names=False, + ).to_dict() + assert missing["trips"] == [None] + + +# --------------------------------------------------------------------------- +# Integration tests that exercise the store directly (no CliRunner/apply). +# --------------------------------------------------------------------------- + + +def _integration_repo_config( + aerospike_online_store_config: dict, collection_suffix: str +) -> RepoConfig: + """Build a RepoConfig targeting the live container with an isolated set name. + + Each test passes a unique ``collection_suffix`` so the module-scoped + Aerospike container can host multiple tests without cross-contamination. + """ + return RepoConfig( + project="itest", + provider="local", + registry="/tmp/reg.db", + online_store={ + **aerospike_online_store_config, + "collection_suffix": collection_suffix, + }, + entity_key_serialization_version=3, + ) + + +def _multi_fv_feature_view(name: str) -> SimpleNamespace: + """Minimal FV duck-typed for the store: exposes .name and .features.""" + return SimpleNamespace( + name=name, + features=[ + SimpleNamespace(name="value", dtype=Int64), + ], + ) + + +@_requires_docker +def test_aerospike_cross_fv_map_cdt_upsert(aerospike_online_store_config): + """Writing two feature views for the same entity must not clobber each other. + + The store uses Aerospike Map CDT ops (``map_put_items``) rather than a + full-record ``put``, so two feature views sharing an entity key live in + separate slots of the ``features`` / ``event_ts`` maps. This test + exercises that guarantee end to end against a real server. + """ + config = _integration_repo_config(aerospike_online_store_config, "cross_fv") + store = AerospikeOnlineStore() + + fv_a = _multi_fv_feature_view("driver_stats") + fv_b = _multi_fv_feature_view("driver_geo") + ek = _entity_key("driver_id", 101) + ts = _utc_now() + + store.online_write_batch( + config, + fv_a, + [(ek, {"value": ValueProto(int64_val=42)}, ts, ts)], + progress=None, + ) + # Second write targets the SAME entity key under a different feature + # view — it must add a new map entry rather than overwrite fv_a. + store.online_write_batch( + config, + fv_b, + [(ek, {"value": ValueProto(int64_val=7)}, ts, ts)], + progress=None, + ) + + results_a = store.online_read(config, fv_a, [ek]) + results_b = store.online_read(config, fv_b, [ek]) + + assert len(results_a) == 1 and len(results_b) == 1 + _, feats_a = results_a[0] + _, feats_b = results_b[0] + assert feats_a is not None, "driver_stats was clobbered by the driver_geo write" + assert feats_b is not None, "driver_geo write did not produce a readable slot" + assert feats_a["value"].int64_val == 42 + assert feats_b["value"].int64_val == 7 + + # Sanity-check the raw record too: both feature views should coexist in + # the ``features`` Map CDT under their own names. + client = store._get_client(config) + _, _, bins = client.get(store._aerospike_key(config, ek)) + assert set(bins["features"].keys()) == {"driver_stats", "driver_geo"} + assert set(bins["event_ts"].keys()) == {"driver_stats", "driver_geo"} + + store.teardown(config, [], []) + + +@_requires_docker +def test_aerospike_update_strips_dropped_feature_view(aerospike_online_store_config): + """``update(tables_to_delete=[...])`` removes a FV's slot from every record. + + The store issues a single background scan that applies + ``map_remove_by_key`` to the ``features`` and ``event_ts`` bins for each + dropped feature view. We verify: + + * The dropped FV's slot disappears from the record (read returns + ``(None, None)``). + * The kept FV is untouched (still readable with its original values). + * The underlying record itself still exists (background scan strips map + entries, not whole records). + """ + config = _integration_repo_config(aerospike_online_store_config, "update") + store = AerospikeOnlineStore() + + fv_keep = _multi_fv_feature_view("keep_fv") + fv_drop = _multi_fv_feature_view("drop_fv") + ek = _entity_key("driver_id", 202) + ts = _utc_now() + + store.online_write_batch( + config, + fv_keep, + [(ek, {"value": ValueProto(int64_val=1)}, ts, ts)], + progress=None, + ) + store.online_write_batch( + config, + fv_drop, + [(ek, {"value": ValueProto(int64_val=999)}, ts, ts)], + progress=None, + ) + + # Pre-condition: both feature views readable. + assert store.online_read(config, fv_keep, [ek])[0][1] is not None + assert store.online_read(config, fv_drop, [ek])[0][1] is not None + + store.update( + config=config, + tables_to_delete=[fv_drop], + tables_to_keep=[fv_keep], + entities_to_delete=[], + entities_to_keep=[], + partial=False, + ) + + # The background scan is asynchronous server-side. Poll up to ~10s for + # the drop_fv slot to disappear before failing — this matches how a real + # Feast caller would experience the post-apply propagation delay. + deadline = time.monotonic() + 10.0 + drop_result = store.online_read(config, fv_drop, [ek])[0] + while drop_result != (None, None) and time.monotonic() < deadline: + time.sleep(0.1) + drop_result = store.online_read(config, fv_drop, [ek])[0] + assert drop_result == (None, None), ( + f"drop_fv slot was not cleared by background scan within 10s; " + f"last result: {drop_result!r}" + ) + + # keep_fv must still be present with its original value. + keep_ts, keep_feats = store.online_read(config, fv_keep, [ek])[0] + assert keep_feats is not None, "keep_fv was removed by the scan" + assert keep_feats["value"].int64_val == 1 + assert keep_ts is not None + + # The record itself still exists — only the dropped FV's map entries + # were removed. Verify directly via a raw get so we notice if a future + # change accidentally escalates to a full-record delete. + client = store._get_client(config) + _, _, bins = client.get(store._aerospike_key(config, ek)) + assert "drop_fv" not in bins["features"] + assert "drop_fv" not in bins["event_ts"] + assert "keep_fv" in bins["features"] + + store.teardown(config, [], []) diff --git a/sdk/python/tests/universal/feature_repos/universal/online_store/aerospike.py b/sdk/python/tests/universal/feature_repos/universal/online_store/aerospike.py new file mode 100644 index 00000000000..de28ad15163 --- /dev/null +++ b/sdk/python/tests/universal/feature_repos/universal/online_store/aerospike.py @@ -0,0 +1,58 @@ +# +# Copyright 2019 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import Any, Dict + +from testcontainers.core.container import DockerContainer +from testcontainers.core.waiting_utils import wait_for_logs + +from tests.universal.feature_repos.universal.online_store_creator import ( + OnlineStoreCreator, +) + +# Aerospike Community Edition — pin a known-good server image so CI stays +# deterministic; bump when Feast's Aerospike server baseline moves. +AEROSPIKE_CE_IMAGE = "aerospike/aerospike-server:8.0.0.10_1" + +# The Aerospike server prints this line once the default "test" namespace has +# finished starting up and the service is accepting client connections. +_READY_LOG_MARKER = "migrations: complete" + + +class AerospikeOnlineStoreCreator(OnlineStoreCreator): + """Spin up a single-node Aerospike CE cluster for universal online-store tests. + + The CE image ships a pre-configured namespace called ``test``, which we + reuse as the Feast namespace (avoiding the need to mount a custom + ``aerospike.conf``). + """ + + def __init__(self, project_name: str, **kwargs): + super().__init__(project_name) + self.container = DockerContainer(AEROSPIKE_CE_IMAGE).with_exposed_ports("3000") + + def create_online_store(self) -> Dict[str, Any]: + self.container.start() + wait_for_logs(container=self.container, predicate=_READY_LOG_MARKER, timeout=60) + exposed_port = int(self.container.get_exposed_port("3000")) + return { + "type": "aerospike", + "hosts": [("127.0.0.1", exposed_port)], + "namespace": "test", + } + + def teardown(self): + self.container.stop() From 267abdfd2a6eb00a44336c71633b7ba1f71fa597 Mon Sep 17 00:00:00 2001 From: Aniket Paluskar Date: Thu, 2 Jul 2026 17:20:45 +0530 Subject: [PATCH 11/16] fix: add debug logging for FIPS mode detection fallback Signed-off-by: Aniket Paluskar --- sdk/python/feast/offline_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/python/feast/offline_server.py b/sdk/python/feast/offline_server.py index cdf3d9edaa6..e130544445f 100644 --- a/sdk/python/feast/offline_server.py +++ b/sdk/python/feast/offline_server.py @@ -55,6 +55,7 @@ def _is_fips_enabled() -> bool: with open("/proc/sys/crypto/fips_enabled") as f: return f.read().strip() == "1" except (FileNotFoundError, PermissionError, OSError): + logger.debug("Could not detect FIPS mode (Linux-only feature)") return False From 4193c0d81624fb542868736db5f9366a248eb381 Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Thu, 2 Jul 2026 20:07:33 -0700 Subject: [PATCH 12/16] fix: Unblock nightly UI build (#6570) * fix: Unblock nightly UI build Signed-off-by: Francisco Javier Arceo * ci: Add UI production build check Signed-off-by: Francisco Javier Arceo --------- Signed-off-by: Francisco Javier Arceo --- .github/workflows/unit_tests.yml | 3 +++ ui/src/components/DataSourceFormModal.tsx | 2 +- ui/src/components/LineageEventsList.tsx | 1 - ui/src/pages/saved-data-sets/CreateDatasetForm.tsx | 2 +- ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx | 2 +- ui/src/pages/saved-data-sets/DatasetUsageTab.tsx | 1 - 6 files changed, 6 insertions(+), 5 deletions(-) diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index 70a99c0ed05..935051a044e 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -103,6 +103,9 @@ jobs: - name: Build yarn rollup working-directory: ./ui run: yarn build:lib + - name: Build production UI + working-directory: ./ui + run: CI=true npm run build --omit=dev - name: Run yarn tests working-directory: ./ui run: yarn test --watchAll=false diff --git a/ui/src/components/DataSourceFormModal.tsx b/ui/src/components/DataSourceFormModal.tsx index 328c614d3d2..ca95bb6aa0e 100644 --- a/ui/src/components/DataSourceFormModal.tsx +++ b/ui/src/components/DataSourceFormModal.tsx @@ -1,4 +1,4 @@ -import React, { useState, useEffect } from "react"; +import React, { useState } from "react"; import { EuiFormRow, EuiFieldText, diff --git a/ui/src/components/LineageEventsList.tsx b/ui/src/components/LineageEventsList.tsx index af9c2ce709d..1c029b9894d 100644 --- a/ui/src/components/LineageEventsList.tsx +++ b/ui/src/components/LineageEventsList.tsx @@ -10,7 +10,6 @@ import { EuiFieldSearch, EuiFlexGroup, EuiFlexItem, - EuiSelect, EuiFlyout, EuiFlyoutHeader, EuiFlyoutBody, diff --git a/ui/src/pages/saved-data-sets/CreateDatasetForm.tsx b/ui/src/pages/saved-data-sets/CreateDatasetForm.tsx index dca9ff3d026..a3be0bae8df 100644 --- a/ui/src/pages/saved-data-sets/CreateDatasetForm.tsx +++ b/ui/src/pages/saved-data-sets/CreateDatasetForm.tsx @@ -106,7 +106,7 @@ const CreateDatasetForm = ({ onClose }: CreateDatasetFormProps) => { const [storageType, setStorageType] = useState("file"); const [storagePath, setStoragePath] = useState(""); const [tags, setTags] = useState([]); - const [allowOverwrite, setAllowOverwrite] = useState(false); + const [allowOverwrite] = useState(false); // Job tracking const [jobId, setJobId] = useState(null); diff --git a/ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx b/ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx index 905227f5327..b765da70241 100644 --- a/ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx +++ b/ui/src/pages/saved-data-sets/DatasetOverviewTab.tsx @@ -64,7 +64,7 @@ const DatasetOverviewTab = () => { ); } - const { isLoading, isSuccess, isError, data } = useLoadDataset(datasetName); + const { isLoading, isError, data } = useLoadDataset(datasetName); if (isLoading) { return ( diff --git a/ui/src/pages/saved-data-sets/DatasetUsageTab.tsx b/ui/src/pages/saved-data-sets/DatasetUsageTab.tsx index 65a7b97fc68..01a9778213e 100644 --- a/ui/src/pages/saved-data-sets/DatasetUsageTab.tsx +++ b/ui/src/pages/saved-data-sets/DatasetUsageTab.tsx @@ -5,7 +5,6 @@ import { EuiText, EuiSpacer, EuiCodeBlock, - EuiHorizontalRule, EuiFlexGroup, EuiFlexItem, EuiCallOut, From 386599e874d4169ca0ac79d60c71bc9448f3c2f1 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Fri, 3 Jul 2026 12:32:30 +0530 Subject: [PATCH 13/16] feat: Implement Iceberg REST Catalog-backed materialization (L3) Signed-off-by: Abhishek Shinde --- .../infra/compute_engines/local/nodes.py | 20 +-- .../infra/compute_engines/spark/nodes.py | 20 +-- .../offline_stores/iceberg/registration.py | 95 +++++++++++- .../test_uc_registration.py | 144 +++++++++++++++++- 4 files changed, 259 insertions(+), 20 deletions(-) diff --git a/sdk/python/feast/infra/compute_engines/local/nodes.py b/sdk/python/feast/infra/compute_engines/local/nodes.py index 2e37863db4b..9878c8e1b32 100644 --- a/sdk/python/feast/infra/compute_engines/local/nodes.py +++ b/sdk/python/feast/infra/compute_engines/local/nodes.py @@ -408,15 +408,17 @@ def execute(self, context: ExecutionContext) -> ArrowTableValue: ) # UC-backed materialization hook (Phase L3) - from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( - write_uc_materialized_data, - ) + catalog_config = getattr(context.repo_config.offline_store, "catalog", None) + if catalog_config is not None: + from feast.infra.offline_stores.iceberg.registration import ( + write_uc_materialized_data, + ) - write_uc_materialized_data( - config=context.repo_config, - fv=self.feature_view, - df=input_table, - project=context.repo_config.project, - ) + write_uc_materialized_data( + catalog_config=catalog_config, + fv=self.feature_view, + df=input_table, + project=context.repo_config.project, + ) return output diff --git a/sdk/python/feast/infra/compute_engines/spark/nodes.py b/sdk/python/feast/infra/compute_engines/spark/nodes.py index 0df3b774973..3e6ced63871 100644 --- a/sdk/python/feast/infra/compute_engines/spark/nodes.py +++ b/sdk/python/feast/infra/compute_engines/spark/nodes.py @@ -591,16 +591,18 @@ def execute(self, context: ExecutionContext) -> DAGValue: spark_df.write.format(file_format).mode("append").save(dest_path) # UC-backed materialization hook (Phase L3) - from feast.infra.offline_stores.contrib.spark_offline_store.uc_registration import ( - write_uc_materialized_data, - ) + catalog_config = getattr(context.repo_config.offline_store, "catalog", None) + if catalog_config is not None: + from feast.infra.offline_stores.iceberg.registration import ( + write_uc_materialized_data, + ) - write_uc_materialized_data( - config=context.repo_config, - fv=self.feature_view, - df=spark_df, - project=context.repo_config.project, - ) + write_uc_materialized_data( + catalog_config=catalog_config, + fv=self.feature_view, + df=spark_df, + project=context.repo_config.project, + ) return DAGValue( data=spark_df, diff --git a/sdk/python/feast/infra/offline_stores/iceberg/registration.py b/sdk/python/feast/infra/offline_stores/iceberg/registration.py index 72b7c2927e4..40bc14be2c6 100644 --- a/sdk/python/feast/infra/offline_stores/iceberg/registration.py +++ b/sdk/python/feast/infra/offline_stores/iceberg/registration.py @@ -17,7 +17,7 @@ """ import logging -from typing import Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import click @@ -28,6 +28,7 @@ load_catalog, table_exists, update_table_properties, + write_materialized_data, ) logger = logging.getLogger(__name__) @@ -186,3 +187,95 @@ def register_uc_feature_tables( f" ✗ Failed to register catalog table: {fv.name} " f"(check logs for details)" ) + + +_MATERIALIZE_OFFLINE_KEY = "uc.materialize_offline" + + +def write_uc_materialized_data( + catalog_config: IcebergCatalogConfig, + fv: FeatureView, + df: Any, + project: str, +) -> None: + """Write materialized features into an Iceberg REST Catalog table. + + Converts ``pa.Table`` to Spark DataFrame if needed, then performs + a MERGE INTO via the catalog manager. + + Skips silently when any guard condition fails: + - ``register_as_feature_table`` tag is ``"false"`` + - ``materialize_offline`` tag is ``"false"`` + - The UC path (catalog/schema) cannot be resolved + - The catalog cannot be reached (logged) + """ + if not _should_register(fv): + logger.info("Skipping UC materialization for %s (opt-out)", fv.name) + return + + if fv.tags.get(_MATERIALIZE_OFFLINE_KEY, "true").lower() == "false": + logger.info( + "Skipping UC materialization for %s (materialize_offline=false)", + fv.name, + ) + return + + catalog_name, schema, table = _resolve_uc_path( + fv, catalog_config.warehouse, catalog_config.namespace + ) + if not catalog_name or not schema: + logger.warning( + "Cannot materialize to UC for %s: missing catalog or schema.", + fv.name, + ) + return + + try: + catalog = load_catalog(catalog_config) + except Exception as e: + logger.warning("Could not create Iceberg catalog for materialization: %s", e) + return + + if not table_exists(catalog, schema, table): + try: + create_feature_table( + catalog=catalog, + namespace=schema, + table_name=table, + entity_columns=fv.entity_columns, + feature_columns=fv.features, + primary_keys=_get_primary_keys(fv), + timestamp_field=getattr(fv.batch_source, "timestamp_field", None), + description=fv.description or "", + owner=fv.owner or "", + project=project, + tags=_build_tags(fv, project), + ) + logger.info("Created catalog table %s.%s.%s", catalog_name, schema, table) + except Exception: + logger.exception( + "Failed to create catalog table %s.%s.%s", + catalog_name, + schema, + table, + ) + raise + + import pyarrow as pa + + if isinstance(df, pa.Table): + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + spark_df = spark.createDataFrame(df.to_pandas()) + else: + spark_df = df + + write_materialized_data( + catalog=catalog, + namespace=schema, + table_name=table, + spark_df=spark_df, + mode="merge", + ) + logger.info("Materialized features to %s.%s.%s", catalog_name, schema, table) diff --git a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py index 0c72a3eac34..02d771fc996 100644 --- a/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py @@ -1,3 +1,4 @@ +from typing import Dict, Optional from unittest.mock import MagicMock, patch from feast import FeatureView @@ -11,7 +12,6 @@ register_uc_feature_tables, write_uc_materialized_data, ) -from feast.repo_config import RepoConfig def make_mock_field(name: str, dtype_str: str = "", nullable: bool = True): @@ -324,3 +324,145 @@ def test_register_uc_feature_tables_skips_when_missing_catalog( mock_create.assert_not_called() mock_update.assert_not_called() + + +# ────────────────────────────────────────────── +# Tests for write_uc_materialized_data (L3) +# ────────────────────────────────────────────── + + +def _make_catalog_config_for_mat( + warehouse: str = "cat", + namespace: str = "sch", +) -> IcebergCatalogConfig: + return IcebergCatalogConfig( + type="rest", + endpoint="https://example.com/api/iceberg", + warehouse=warehouse, + namespace=namespace, + ) + + +def _make_mock_fv_for_mat( + name: str = "test_fv", + tags: Optional[Dict] = None, +) -> FeatureView: + fv = MagicMock(spec=FeatureView) + fv.name = name + fv.tags = tags or {} + fv.entity_columns = [] + fv.features = [] + fv.batch_source = MagicMock() + fv.batch_source.timestamp_field = None + fv.description = "" + fv.owner = None + return fv + + +def test_write_uc_materialized_data_skips_opt_out(): + """Tag uc.register_as_feature_table=false → skip.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat(tags={"uc.register_as_feature_table": "false"}) + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_load.assert_not_called() + + +def test_write_uc_materialized_data_skips_materialize_offline_false(): + """Tag uc.materialize_offline=false → skip.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat(tags={"uc.materialize_offline": "false"}) + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_load.assert_not_called() + + +def test_write_uc_materialized_data_skips_missing_catalog(): + """Missing catalog/schema → skip (log warning).""" + catalog_config = _make_catalog_config_for_mat(warehouse="", namespace="") + fv = _make_mock_fv_for_mat() + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_load.assert_not_called() + + +def test_write_uc_materialized_data_skips_catalog_load_failure(): + """load_catalog raises → skip.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat() + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + side_effect=Exception("connection error"), + ) as mock_load: + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_load.assert_called_once() + + +def test_write_uc_materialized_data_writes_merge_existing_table(): + """Existing table → write_materialized_data (merge) is called.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat() + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + with patch( + "feast.infra.offline_stores.iceberg.registration.table_exists", + return_value=True, + ) as mock_exists: + with patch( + "feast.infra.offline_stores.iceberg.registration.write_materialized_data", + ) as mock_write: + mock_catalog = MagicMock() + mock_load.return_value = mock_catalog + + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_exists.assert_called_once() + mock_write.assert_called_once() + _, kwargs = mock_write.call_args + assert kwargs["namespace"] == "sch" + assert kwargs["table_name"] == "test_fv" + assert kwargs["mode"] == "merge" + + +def test_write_uc_materialized_data_creates_then_writes(): + """New table → create_feature_table then write_materialized_data.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat() + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + with patch( + "feast.infra.offline_stores.iceberg.registration.table_exists", + return_value=False, + ): + with patch( + "feast.infra.offline_stores.iceberg.registration.create_feature_table", + ) as mock_create: + with patch( + "feast.infra.offline_stores.iceberg.registration.write_materialized_data", + ) as mock_write: + mock_catalog = MagicMock() + mock_load.return_value = mock_catalog + + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_create.assert_called_once() + mock_write.assert_called_once() + assert mock_write.call_args[1]["mode"] == "merge" From 40a3de3bf586ab2ede49070704231adca0281dc8 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sat, 4 Jul 2026 08:18:32 +0530 Subject: [PATCH 14/16] fix: Resolve 4 mypy errors in spark offline store, iceberg catalog, and spark utils Signed-off-by: Abhishek Shinde --- .../offline_stores/contrib/spark_offline_store/spark.py | 7 ++++--- .../offline_stores/contrib/spark_offline_store/utils.py | 2 +- .../feast/infra/offline_stores/iceberg/catalog_manager.py | 4 ++-- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py index f7030417cbb..ec2a036c237 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py @@ -195,9 +195,10 @@ def get_historical_features( RuntimeWarning, ) - spark_session = _resolve_spark_session_for_source( - feature_views[0].batch_source, config - ) + batch_source = feature_views[0].batch_source + if batch_source is None: + raise ValueError("batch_source is required for spark offline store") + spark_session = _resolve_spark_session_for_source(batch_source, config) tmp_entity_df_table_name = offline_utils.get_temp_entity_table_name() # Non-entity mode: synthesize a left table and timestamp range from start/end dates to avoid requiring entity_df. diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py index b797bd2a6c9..8c0d2490116 100644 --- a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py @@ -43,7 +43,7 @@ def get_databricks_connect_session( spark_session = SparkSession.getActiveSession() if spark_session and host: try: - current_uri = spark_session.conf.get("spark.connect.client.uri", "") + current_uri = spark_session.conf.get("spark.connect.client.uri", "") or "" expected_uri = _expected_connect_uri(host) if expected_uri not in current_uri: logger.info( diff --git a/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py b/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py index 1d4129595bb..c6a7e79c05c 100644 --- a/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py +++ b/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py @@ -350,7 +350,7 @@ def _merge_into( pk_cols = [pk.strip() for pk in primary_keys_str.split(",") if pk.strip()] if not pk_cols: - spark_df.sql_ctx.sql(f"INSERT INTO {full_name} SELECT * FROM {temp_view}") + spark_df.sparkSession.sql(f"INSERT INTO {full_name} SELECT * FROM {temp_view}") return join_cond = " AND ".join(f"target.`{pk}` = source.`{pk}`" for pk in pk_cols) @@ -368,4 +368,4 @@ def _merge_into( f"WHEN NOT MATCHED THEN INSERT ({insert_cols}) VALUES ({source_cols})" ) - spark_df.sql_ctx.sql(merge_sql) + spark_df.sparkSession.sql(merge_sql) From a52ea37feae72c9150534ae2e3ddb5cff7725a8f Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sat, 4 Jul 2026 08:53:04 +0530 Subject: [PATCH 15/16] fix: Eliminate TOCTOU race condition in Iceberg table registration Signed-off-by: Abhishek Shinde --- .../offline_stores/iceberg/registration.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/sdk/python/feast/infra/offline_stores/iceberg/registration.py b/sdk/python/feast/infra/offline_stores/iceberg/registration.py index 72b7c2927e4..750fa543c2c 100644 --- a/sdk/python/feast/infra/offline_stores/iceberg/registration.py +++ b/sdk/python/feast/infra/offline_stores/iceberg/registration.py @@ -26,7 +26,6 @@ from feast.infra.offline_stores.iceberg.catalog_manager import ( create_feature_table, load_catalog, - table_exists, update_table_properties, ) @@ -107,35 +106,40 @@ def _register_single_feature_view( owner = fv.owner or "" tags = _build_tags(fv, project) - if table_exists(catalog, namespace, table_name): - update_table_properties( + try: + create_feature_table( catalog=catalog, namespace=namespace, table_name=table_name, + entity_columns=fv.entity_columns, + feature_columns=fv.features, + primary_keys=primary_keys, + timestamp_field=timestamp_field, description=description, owner=owner, project=project, tags=tags, ) click.echo( - f" ✓ Updated catalog table: {catalog_name}.{namespace}.{table_name}" + f" ✓ Created catalog table: {catalog_name}.{namespace}.{table_name}" ) - else: - create_feature_table( + except Exception as e: + from pyiceberg.exceptions import TableAlreadyExistsError + + if not isinstance(e, TableAlreadyExistsError): + raise + + update_table_properties( catalog=catalog, namespace=namespace, table_name=table_name, - entity_columns=fv.entity_columns, - feature_columns=fv.features, - primary_keys=primary_keys, - timestamp_field=timestamp_field, description=description, owner=owner, project=project, tags=tags, ) click.echo( - f" ✓ Created catalog table: {catalog_name}.{namespace}.{table_name}" + f" ✓ Updated catalog table: {catalog_name}.{namespace}.{table_name}" ) From 23d8c15d827cd677c4897dbd683047bf81edfdcb Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sat, 4 Jul 2026 09:56:44 +0530 Subject: [PATCH 16/16] fix: Resolve materialization issues Signed-off-by: Abhishek Shinde --- .../infra/compute_engines/local/nodes.py | 6 ++- .../infra/compute_engines/spark/nodes.py | 6 ++- .../offline_stores/iceberg/catalog_manager.py | 9 ++-- .../offline_stores/iceberg/registration.py | 49 +++++++++++-------- 4 files changed, 45 insertions(+), 25 deletions(-) diff --git a/sdk/python/feast/infra/compute_engines/local/nodes.py b/sdk/python/feast/infra/compute_engines/local/nodes.py index 9878c8e1b32..4ede3896a7e 100644 --- a/sdk/python/feast/infra/compute_engines/local/nodes.py +++ b/sdk/python/feast/infra/compute_engines/local/nodes.py @@ -409,7 +409,11 @@ def execute(self, context: ExecutionContext) -> ArrowTableValue: # UC-backed materialization hook (Phase L3) catalog_config = getattr(context.repo_config.offline_store, "catalog", None) - if catalog_config is not None: + from feast.infra.offline_stores.iceberg.catalog_config import ( + IcebergCatalogConfig, + ) + + if isinstance(catalog_config, IcebergCatalogConfig): from feast.infra.offline_stores.iceberg.registration import ( write_uc_materialized_data, ) diff --git a/sdk/python/feast/infra/compute_engines/spark/nodes.py b/sdk/python/feast/infra/compute_engines/spark/nodes.py index 3e6ced63871..50c8256a20f 100644 --- a/sdk/python/feast/infra/compute_engines/spark/nodes.py +++ b/sdk/python/feast/infra/compute_engines/spark/nodes.py @@ -592,7 +592,11 @@ def execute(self, context: ExecutionContext) -> DAGValue: # UC-backed materialization hook (Phase L3) catalog_config = getattr(context.repo_config.offline_store, "catalog", None) - if catalog_config is not None: + from feast.infra.offline_stores.iceberg.catalog_config import ( + IcebergCatalogConfig, + ) + + if isinstance(catalog_config, IcebergCatalogConfig): from feast.infra.offline_stores.iceberg.registration import ( write_uc_materialized_data, ) diff --git a/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py b/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py index c6a7e79c05c..9e4ea9fd50a 100644 --- a/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py +++ b/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py @@ -285,9 +285,12 @@ def update_table_properties( """Update Iceberg table properties with current Feast metadata.""" table = catalog.load_table((namespace, table_name)) + existing_pks_str = table.properties.get(FEAST_PRIMARY_KEYS_PROP, "") + existing_ts = table.properties.get(FEAST_TIMESTAMP_KEY_PROP, "") + new_props = _build_iceberg_properties( - primary_keys=[], - timestamp_key=None, + primary_keys=[pk.strip() for pk in existing_pks_str.split(",") if pk.strip()], + timestamp_key=existing_ts or None, description=description, owner=owner, project=project, @@ -340,7 +343,7 @@ def _merge_into( Requires Spark 3.x with Iceberg support. """ - full_name = f"{namespace}.{table_name}" + full_name = f"`{namespace}`.`{table_name}`" temp_view = f"_feast_merge_{table_name}" spark_df.createOrReplaceTempView(temp_view) diff --git a/sdk/python/feast/infra/offline_stores/iceberg/registration.py b/sdk/python/feast/infra/offline_stores/iceberg/registration.py index 60c33c6ac88..318bb8027ec 100644 --- a/sdk/python/feast/infra/offline_stores/iceberg/registration.py +++ b/sdk/python/feast/infra/offline_stores/iceberg/registration.py @@ -240,23 +240,25 @@ def write_uc_materialized_data( logger.warning("Could not create Iceberg catalog for materialization: %s", e) return - if not table_exists(catalog, schema, table): - try: - create_feature_table( - catalog=catalog, - namespace=schema, - table_name=table, - entity_columns=fv.entity_columns, - feature_columns=fv.features, - primary_keys=_get_primary_keys(fv), - timestamp_field=getattr(fv.batch_source, "timestamp_field", None), - description=fv.description or "", - owner=fv.owner or "", - project=project, - tags=_build_tags(fv, project), - ) - logger.info("Created catalog table %s.%s.%s", catalog_name, schema, table) - except Exception: + try: + create_feature_table( + catalog=catalog, + namespace=schema, + table_name=table, + entity_columns=fv.entity_columns, + feature_columns=fv.features, + primary_keys=_get_primary_keys(fv), + timestamp_field=getattr(fv.batch_source, "timestamp_field", None), + description=fv.description or "", + owner=fv.owner or "", + project=project, + tags=_build_tags(fv, project), + ) + logger.info("Created catalog table %s.%s.%s", catalog_name, schema, table) + except Exception as e: + from pyiceberg.exceptions import TableAlreadyExistsError + + if not isinstance(e, TableAlreadyExistsError): logger.exception( "Failed to create catalog table %s.%s.%s", catalog_name, @@ -268,10 +270,17 @@ def write_uc_materialized_data( import pyarrow as pa if isinstance(df, pa.Table): - from pyspark.sql import SparkSession + try: + from pyspark.sql import SparkSession - spark = SparkSession.builder.getOrCreate() - spark_df = spark.createDataFrame(df.to_pandas()) + spark = SparkSession.builder.getOrCreate() + spark_df = spark.createDataFrame(df) + except ImportError: + logger.warning( + "pyspark is not installed; skipping UC materialization for %s", + fv.name, + ) + return else: spark_df = df