diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index 75f406c8e2c..701027d7537 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.iceberg.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.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). + + 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.iceberg.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.catalog, 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/compute_engines/local/nodes.py b/sdk/python/feast/infra/compute_engines/local/nodes.py index 9d3e1a48881..4ede3896a7e 100644 --- a/sdk/python/feast/infra/compute_engines/local/nodes.py +++ b/sdk/python/feast/infra/compute_engines/local/nodes.py @@ -407,4 +407,22 @@ def execute(self, context: ExecutionContext) -> ArrowTableValue: progress=lambda x: None, ) + # UC-backed materialization hook (Phase L3) + catalog_config = getattr(context.repo_config.offline_store, "catalog", None) + from feast.infra.offline_stores.iceberg.catalog_config import ( + IcebergCatalogConfig, + ) + + if isinstance(catalog_config, IcebergCatalogConfig): + from feast.infra.offline_stores.iceberg.registration import ( + write_uc_materialized_data, + ) + + write_uc_materialized_data( + catalog_config=catalog_config, + fv=self.feature_view, + df=input_table, + project=context.repo_config.project, + ) + return output diff --git a/sdk/python/feast/infra/compute_engines/spark/nodes.py b/sdk/python/feast/infra/compute_engines/spark/nodes.py index 92964b72bc9..50c8256a20f 100644 --- a/sdk/python/feast/infra/compute_engines/spark/nodes.py +++ b/sdk/python/feast/infra/compute_engines/spark/nodes.py @@ -590,6 +590,24 @@ def execute(self, context: ExecutionContext) -> DAGValue: ) spark_df.write.format(file_format).mode("append").save(dest_path) + # UC-backed materialization hook (Phase L3) + catalog_config = getattr(context.repo_config.offline_store, "catalog", None) + from feast.infra.offline_stores.iceberg.catalog_config import ( + IcebergCatalogConfig, + ) + + if isinstance(catalog_config, IcebergCatalogConfig): + from feast.infra.offline_stores.iceberg.registration import ( + write_uc_materialized_data, + ) + + write_uc_materialized_data( + catalog_config=catalog_config, + fv=self.feature_view, + df=spark_df, + project=context.repo_config.project, + ) + return DAGValue( data=spark_df, format=DAGFormat.SPARK, 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..7eacd5a4481 --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/databricks_uc.py @@ -0,0 +1,373 @@ +import logging +from datetime import date, datetime +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union + +import pandas as pd +import pyarrow +import pyspark +from pydantic import StrictBool, 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.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 + +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 catalog feature tables on ``feast apply``.""" + + catalog: Optional[StrictStr] = None + """Default catalog for feature tables. Overrides ``IcebergCatalogConfig.warehouse``.""" + + uc_schema: Optional[StrictStr] = None + """Default schema for feature tables. Overrides ``IcebergCatalogConfig.namespace``.""" + + +class DatabricksUCOfflineStoreConfig(SparkOfflineStoreConfig): + """Offline store configuration for Databricks Unity Catalog. + + 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. + """ + + type: StrictStr = "databricks_uc" + """Offline store type selector""" + + catalog: IcebergCatalogConfig + """Iceberg REST Catalog connection configuration. + + For Databricks UC:: + + 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``.""" + + +def get_databricks_session( + store_config: DatabricksUCOfflineStoreConfig, +) -> SparkSession: + """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: + catalog_cfg = store_config.catalog + endpoint = catalog_cfg.endpoint + + # 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: + conn_str = f"sc://{workspace_host}:443/" + params = [] + if token: + params.append(f"token={token}") + 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 + + 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() + + assert spark_session is not None + + spark_session.conf.set("spark.sql.parser.quotedRegexColumnNames", "true") + + 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, + 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) + 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, + ) + + @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) 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..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 @@ -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,9 +195,10 @@ def get_historical_features( RuntimeWarning, ) - spark_session = get_spark_session_or_start_new_with_repoconfig( - store_config=config.offline_store - ) + 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. @@ -329,8 +337,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 +352,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 +401,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 +457,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 +506,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 +1184,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 +1238,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/utils.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/utils.py new file mode 100644 index 00000000000..8c0d2490116 --- /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", "") or "" + 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..9e4ea9fd50a --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/iceberg/catalog_manager.py @@ -0,0 +1,374 @@ +"""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)) + + existing_pks_str = table.properties.get(FEAST_PRIMARY_KEYS_PROP, "") + existing_ts = table.properties.get(FEAST_TIMESTAMP_KEY_PROP, "") + + new_props = _build_iceberg_properties( + primary_keys=[pk.strip() for pk in existing_pks_str.split(",") if pk.strip()], + timestamp_key=existing_ts or 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.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) + + 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.sparkSession.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..318bb8027ec --- /dev/null +++ b/sdk/python/feast/infra/offline_stores/iceberg/registration.py @@ -0,0 +1,294 @@ +"""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 Any, 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, + update_table_properties, + write_materialized_data, +) + +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) + + 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" ✓ Created catalog table: {catalog_name}.{namespace}.{table_name}" + ) + 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, + description=description, + owner=owner, + project=project, + tags=tags, + ) + click.echo( + f" ✓ Updated 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)" + ) + + +_MATERIALIZE_OFFLINE_KEY = "uc.materialize_offline" + + +def write_uc_materialized_data( + catalog_config: IcebergCatalogConfig, + fv: FeatureView, + df: Any, + project: str, +) -> None: + """Write materialized features into an Iceberg REST Catalog table. + + Converts ``pa.Table`` to Spark DataFrame if needed, then performs + a MERGE INTO via the catalog manager. + + Skips silently when any guard condition fails: + - ``register_as_feature_table`` tag is ``"false"`` + - ``materialize_offline`` tag is ``"false"`` + - The UC path (catalog/schema) cannot be resolved + - The catalog cannot be reached (logged) + """ + if not _should_register(fv): + logger.info("Skipping UC materialization for %s (opt-out)", fv.name) + return + + if fv.tags.get(_MATERIALIZE_OFFLINE_KEY, "true").lower() == "false": + logger.info( + "Skipping UC materialization for %s (materialize_offline=false)", + fv.name, + ) + return + + catalog_name, schema, table = _resolve_uc_path( + fv, catalog_config.warehouse, catalog_config.namespace + ) + if not catalog_name or not schema: + logger.warning( + "Cannot materialize to UC for %s: missing catalog or schema.", + fv.name, + ) + return + + try: + catalog = load_catalog(catalog_config) + except Exception as e: + logger.warning("Could not create Iceberg catalog for materialization: %s", e) + return + + try: + create_feature_table( + catalog=catalog, + namespace=schema, + table_name=table, + entity_columns=fv.entity_columns, + feature_columns=fv.features, + primary_keys=_get_primary_keys(fv), + timestamp_field=getattr(fv.batch_source, "timestamp_field", None), + description=fv.description or "", + owner=fv.owner or "", + project=project, + tags=_build_tags(fv, project), + ) + logger.info("Created catalog table %s.%s.%s", catalog_name, schema, table) + except Exception as e: + from pyiceberg.exceptions import TableAlreadyExistsError + + if not isinstance(e, TableAlreadyExistsError): + logger.exception( + "Failed to create catalog table %s.%s.%s", + catalog_name, + schema, + table, + ) + raise + + import pyarrow as pa + + if isinstance(df, pa.Table): + try: + from pyspark.sql import SparkSession + + spark = SparkSession.builder.getOrCreate() + spark_df = spark.createDataFrame(df) + except ImportError: + logger.warning( + "pyspark is not installed; skipping UC materialization for %s", + fv.name, + ) + return + else: + spark_df = df + + write_materialized_data( + catalog=catalog, + namespace=schema, + table_name=table, + spark_df=spark_df, + mode="merge", + ) + logger.info("Materialized features to %s.%s.%s", catalog_name, schema, table) diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index 971c0325f4b..847f933d4f6 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -96,6 +96,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..73a6506c7cc --- /dev/null +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_databricks_uc.py @@ -0,0 +1,233 @@ +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.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", + "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.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", + catalog=_catalog_config(), + 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", + catalog=_catalog_config(warehouse="my_catalog", namespace="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", + 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"}, + ) + + session = get_databricks_session(config) + + assert session == mock_session + mock_builder.remote.assert_called_once_with( + "sc://adb-12345.azuredatabricks.net:443/" + ) + + +@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", + catalog=_catalog_config( + endpoint="https://adb-123.databricks.com/api/2.1/unity-catalog/iceberg", + ), + ), + ) + + 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", + catalog=_catalog_config( + endpoint="https://adb-123.databricks.com/api/2.1/unity-catalog/iceberg", + ), + ), + ) + + 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, + ) 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 new file mode 100644 index 00000000000..02d771fc996 --- /dev/null +++ b/sdk/python/tests/unit/infra/offline_stores/contrib/spark_offline_store/test_uc_registration.py @@ -0,0 +1,468 @@ +from typing import Dict, Optional +from unittest.mock import MagicMock, patch + +from feast import FeatureView +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, + register_uc_feature_tables, + write_uc_materialized_data, +) + + +def make_mock_field(name: str, dtype_str: str = "", nullable: bool = True): + field = MagicMock() + field.name = name + if dtype_str: + + class FakeDtype: + def __str__(self): + return dtype_str + + field.dtype = FakeDtype() + else: + field.dtype = None + return field + + +def test_feast_to_iceberg_type_none(): + assert _feast_to_iceberg_type(None) == "string" + + +def test_feast_to_iceberg_type_int32(): + assert _feast_to_iceberg_type(make_mock_field("col", "Int32").dtype) == "int" + + +def test_feast_to_iceberg_type_int64(): + assert _feast_to_iceberg_type(make_mock_field("col", "Int64").dtype) == "long" + + +def test_feast_to_iceberg_type_float32(): + assert _feast_to_iceberg_type(make_mock_field("col", "Float32").dtype) == "float" + + +def test_feast_to_iceberg_type_float64(): + assert _feast_to_iceberg_type(make_mock_field("col", "Float64").dtype) == "double" + + +def test_feast_to_iceberg_type_string(): + assert _feast_to_iceberg_type(make_mock_field("col", "String").dtype) == "string" + + +def test_feast_to_iceberg_type_bool(): + assert _feast_to_iceberg_type(make_mock_field("col", "Bool").dtype) == "boolean" + + +def test_feast_to_iceberg_type_bytes(): + assert _feast_to_iceberg_type(make_mock_field("col", "Bytes").dtype) == "binary" + + +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_iceberg_type_list(): + t = _feast_to_iceberg_type(make_mock_field("col", "List").dtype) + assert t == "list" + + +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_iceberg_type_unknown(): + t = _feast_to_iceberg_type(make_mock_field("col", "UnknownType").dtype) + assert t == "string" + + +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_tags(): + fv = MagicMock(spec=FeatureView) + fv.tags = { + "env": "prod", + "uc.register_as_feature_table": "false", + "owner": "team_a", + } + tags = _build_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_tags_empty_tags(): + fv = MagicMock(spec=FeatureView) + fv.tags = {} + tags = _build_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 _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) + + +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 + + +@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() + + 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( + "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" + + +@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 + + config = _make_catalog_config() + + 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( + "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") + + 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.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 = _make_catalog_config() + + 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 + + with patch( + "feast.infra.offline_stores.iceberg.registration.create_feature_table" + ) as mock_create: + register_uc_feature_tables(config, [fv], "proj") + + mock_create.assert_not_called() + + +@patch("feast.infra.offline_stores.iceberg.registration.load_catalog") +def test_register_uc_feature_tables_skips_when_missing_catalog( + mock_load_catalog, +): + mock_catalog = MagicMock() + mock_load_catalog.return_value = mock_catalog + + 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 + + # Use a config with empty warehouse/namespace to simulate missing catalog info + config_no_warehouse = _make_catalog_config(warehouse="", namespace="") + + with patch( + "feast.infra.offline_stores.iceberg.registration.create_feature_table" + ) as mock_create: + with patch( + "feast.infra.offline_stores.iceberg.registration.update_table_properties" + ) as mock_update: + register_uc_feature_tables(config_no_warehouse, [fv], "proj") + + mock_create.assert_not_called() + mock_update.assert_not_called() + + +# ────────────────────────────────────────────── +# Tests for write_uc_materialized_data (L3) +# ────────────────────────────────────────────── + + +def _make_catalog_config_for_mat( + warehouse: str = "cat", + namespace: str = "sch", +) -> IcebergCatalogConfig: + return IcebergCatalogConfig( + type="rest", + endpoint="https://example.com/api/iceberg", + warehouse=warehouse, + namespace=namespace, + ) + + +def _make_mock_fv_for_mat( + name: str = "test_fv", + tags: Optional[Dict] = None, +) -> FeatureView: + fv = MagicMock(spec=FeatureView) + fv.name = name + fv.tags = tags or {} + fv.entity_columns = [] + fv.features = [] + fv.batch_source = MagicMock() + fv.batch_source.timestamp_field = None + fv.description = "" + fv.owner = None + return fv + + +def test_write_uc_materialized_data_skips_opt_out(): + """Tag uc.register_as_feature_table=false → skip.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat(tags={"uc.register_as_feature_table": "false"}) + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_load.assert_not_called() + + +def test_write_uc_materialized_data_skips_materialize_offline_false(): + """Tag uc.materialize_offline=false → skip.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat(tags={"uc.materialize_offline": "false"}) + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_load.assert_not_called() + + +def test_write_uc_materialized_data_skips_missing_catalog(): + """Missing catalog/schema → skip (log warning).""" + catalog_config = _make_catalog_config_for_mat(warehouse="", namespace="") + fv = _make_mock_fv_for_mat() + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_load.assert_not_called() + + +def test_write_uc_materialized_data_skips_catalog_load_failure(): + """load_catalog raises → skip.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat() + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + side_effect=Exception("connection error"), + ) as mock_load: + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_load.assert_called_once() + + +def test_write_uc_materialized_data_writes_merge_existing_table(): + """Existing table → write_materialized_data (merge) is called.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat() + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + with patch( + "feast.infra.offline_stores.iceberg.registration.table_exists", + return_value=True, + ) as mock_exists: + with patch( + "feast.infra.offline_stores.iceberg.registration.write_materialized_data", + ) as mock_write: + mock_catalog = MagicMock() + mock_load.return_value = mock_catalog + + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_exists.assert_called_once() + mock_write.assert_called_once() + _, kwargs = mock_write.call_args + assert kwargs["namespace"] == "sch" + assert kwargs["table_name"] == "test_fv" + assert kwargs["mode"] == "merge" + + +def test_write_uc_materialized_data_creates_then_writes(): + """New table → create_feature_table then write_materialized_data.""" + catalog_config = _make_catalog_config_for_mat() + fv = _make_mock_fv_for_mat() + + with patch( + "feast.infra.offline_stores.iceberg.registration.load_catalog", + ) as mock_load: + with patch( + "feast.infra.offline_stores.iceberg.registration.table_exists", + return_value=False, + ): + with patch( + "feast.infra.offline_stores.iceberg.registration.create_feature_table", + ) as mock_create: + with patch( + "feast.infra.offline_stores.iceberg.registration.write_materialized_data", + ) as mock_write: + mock_catalog = MagicMock() + mock_load.return_value = mock_catalog + + write_uc_materialized_data(catalog_config, fv, MagicMock(), "proj") + + mock_create.assert_called_once() + mock_write.assert_called_once() + assert mock_write.call_args[1]["mode"] == "merge"