From 0800f135332adf57b68a7fb95ac2faee9125f77f Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sun, 14 Jun 2026 15:38:45 +0530 Subject: [PATCH 1/8] 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 2/8] 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 3/8] 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 4/8] 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 5/8] 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 6/8] 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 40a3de3bf586ab2ede49070704231adca0281dc8 Mon Sep 17 00:00:00 2001 From: Abhishek Shinde Date: Sat, 4 Jul 2026 08:18:32 +0530 Subject: [PATCH 7/8] 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 8/8] 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}" )