diff --git a/docs/getting-started/components/open-telemetry.md b/docs/getting-started/components/open-telemetry.md index bbad5ed296e..88daf6ed910 100644 --- a/docs/getting-started/components/open-telemetry.md +++ b/docs/getting-started/components/open-telemetry.md @@ -38,7 +38,7 @@ The OpenTelemetry integration in Feast consists of several components working to To add monitoring to the Feast Feature Server, follow these steps: ### 1. Deploy Prometheus Operator -Follow the [Prometheus Operator documentation](https://github.com/prometheus-operator/prometheus-operator/blob/main/Documentation/user-guides/getting-started.md) to install the operator. +Follow the [Prometheus Operator documentation](https://github.com/prometheus-operator/prometheus-operator/blob/main/Documentation/getting-started/installation.md) to install the operator. ### 2. Deploy OpenTelemetry Operator Before installing the OpenTelemetry Operator: diff --git a/docs/getting-started/concepts/feature-retrieval.md b/docs/getting-started/concepts/feature-retrieval.md index 7c4dec08ef9..e8e3b9fa9b1 100644 --- a/docs/getting-started/concepts/feature-retrieval.md +++ b/docs/getting-started/concepts/feature-retrieval.md @@ -15,7 +15,7 @@ Each of these retrieval mechanisms accept: * some way of specifying entities (to fetch features for) * some way to specify the features to fetch (either via [feature services](feature-retrieval.md#feature-services), which group features needed for a model version, or [feature references](feature-retrieval.md#feature-references)) -Before beginning, you need to instantiate a local `FeatureStore` object that knows how to parse the registry (see [more details](https://docs.feast.dev/getting-started/concepts/registry)) +Before beginning, you need to instantiate a local `FeatureStore` object that knows how to parse the registry (see [more details](https://docs.feast.dev/getting-started/components/registry)) For code examples of how the below work, inspect the generated repository from `feast init -t [YOUR TEMPLATE]` (`gcp`, `snowflake`, and `aws` are the most fully fleshed). diff --git a/docs/reference/compute-engine/ray.md b/docs/reference/compute-engine/ray.md index 19604e21788..955d060b987 100644 --- a/docs/reference/compute-engine/ray.md +++ b/docs/reference/compute-engine/ray.md @@ -35,7 +35,7 @@ The Ray compute engine provides: The Ray compute engine follows Feast's DAG-based architecture: ``` -EntityDF → RayReadNode → RayJoinNode → RayFilterNode → RayAggregationNode → RayTransformationNode → Output +EntityDF → RayReadNode → RayTransformationNode → RayJoinNode → RayFilterNode → RayAggregationNode → Output ``` ### Core Components diff --git a/docs/reference/data-sources/iceberg.md b/docs/reference/data-sources/iceberg.md index 6402a7ec419..4d7dbdb7e73 100644 --- a/docs/reference/data-sources/iceberg.md +++ b/docs/reference/data-sources/iceberg.md @@ -84,6 +84,53 @@ my_uc_source = UnityCatalogSource( When `endpoint` is omitted, it defaults to `{DATABRICKS_HOST}/api/2.1/unity-catalog/iceberg`. When `token_env_var` is omitted, it defaults to `DATABRICKS_TOKEN`. +### PolarisSource (Apache Polaris) {#polaris-source} + +Extends `IcebergSource` for [Apache Polaris](https://polaris.apache.org/). `warehouse` is the Polaris +catalog name. When `endpoint` is omitted it defaults to `$POLARIS_HOST/api/catalog`. + +```python +from feast.infra.data_sources.contrib.iceberg_catalog import PolarisSource + +my_polaris_source = PolarisSource( + warehouse="my_polaris_catalog", + namespace="ml", + table="driver_stats", + timestamp_field="event_timestamp", +) +``` + +Authentication is read from the environment (only the variable *names* are stored in the registry): + +1. OAuth2 client credentials: `POLARIS_CLIENT_ID` and `POLARIS_CLIENT_SECRET` (override the names with + `client_id_env_var` / `client_secret_env_var`). Tokens are refreshed automatically before they expire. +2. Otherwise a bearer token from `POLARIS_TOKEN` (override with `token_env_var`). + +To record Feast metadata on the Polaris tables when you run `feast apply`, select the Polaris provider: + +```yaml +# feature_store.yaml +provider: polaris +``` + +For every feature view backed by a `PolarisSource` (with `sync_metadata=True`, the default) the provider sets +`feast.project`, `feast.feature_view`, `feast.owner`, `feast.description`, `feast.tag.` and +`feast.is_feature_table` as table properties through the standard Iceberg REST `updateTable` operation. +When a feature view is deleted the table is kept and marked `feast.is_feature_table=false` with a +`feast.deleted_at` timestamp. A failed metadata sync is logged and never fails `feast apply`. + +Polaris has no lineage API, so `sync_lineage` is not available. + +To create a source for every table in a catalog: + +```python +sources = PolarisSource.discover( + warehouse="my_polaris_catalog", + namespaces=["ml"], # omit to scan all namespaces + timestamp_fields=("event_timestamp",), # tables without such a column are skipped +) +``` + ### Full Feature View Example ```python @@ -145,6 +192,21 @@ driver_stats_fv = FeatureView( | `register_as_feature_table` | `bool` | Register as UC feature table on `feast apply` (default: `True`) | | `sync_lineage` | `bool` | Sync lineage metadata to Unity Catalog (default: `True`) | +### PolarisSource (additional parameters) + +`PolarisSource` always uses `catalog_type="rest"`; `warehouse` is the Polaris catalog name. + +| Parameter | Type | Description | +| :--- | :--- | :--- | +| `endpoint` | `str` | Defaults to `$POLARIS_HOST/api/catalog` | +| `token_env_var` | `str` | Bearer token variable (default: `POLARIS_TOKEN`) | +| `client_id_env_var` | `str` | OAuth client id variable (default: `POLARIS_CLIENT_ID`) | +| `client_secret_env_var` | `str` | OAuth client secret variable (default: `POLARIS_CLIENT_SECRET`) | +| `oauth_scope` | `str` | OAuth scope (default: `PRINCIPAL_ROLE:ALL`) | +| `oauth_server_uri` | `str` | Token endpoint if not `/v1/oauth/tokens`, e.g. an external IdP | +| `credential_vending` | `bool` | Request vended storage credentials for reads/writes (default: `True`) | +| `sync_metadata` | `bool` | Sync Feast metadata to table properties on `feast apply` with the `polaris` provider (default: `True`) | + ## Supported Types | Iceberg Type | Feast Type | diff --git a/docs/reference/data-sources/push.md b/docs/reference/data-sources/push.md index fb7b88a1810..ad69d464f4e 100644 --- a/docs/reference/data-sources/push.md +++ b/docs/reference/data-sources/push.md @@ -84,3 +84,22 @@ streamingDF.writeStream.foreachBatch(feast_writer).start() ``` This can also be used under the hood by a contrib stream processor (see [Tutorial: Building streaming features](../../tutorials/building-streaming-features.md)) + +## External Pipeline Lineage (upstream_feature_view_names) + +When feature transformations occurring outside Feast consume upstream Feast feature views before pushing results into Feast, this dependency metadata can be captured by setting `upstream_feature_view_names` on `PushSource` to visualize the relationship in the Feast Web UI lineage graph: + +```python +push_source2 = PushSource( + name="push_source2", + upstream_feature_view_names=[fv], +) + +fv2 = FeatureView( + name="feature view 2", + entities=[user], + schema=[Field(name="derived_value", dtype=Int64)], + source=push_source2, +) +``` + diff --git a/docs/reference/data-sources/table-formats.md b/docs/reference/data-sources/table-formats.md index 047f1771d18..7d1e85d621c 100644 --- a/docs/reference/data-sources/table-formats.md +++ b/docs/reference/data-sources/table-formats.md @@ -31,6 +31,10 @@ lance_format = LanceFormat( ) ``` +The DuckDB offline store reads a `LanceSource` as a lazy scan, so only the +columns a query selects are read from the dataset. Materializing a few scalar +features does not load the embedding columns stored beside them. + ### Apache Iceberg [Apache Iceberg](https://iceberg.apache.org/) is an open table format designed for huge analytic datasets. It provides: diff --git a/go/infra/docker/feature-server/Dockerfile b/go/infra/docker/feature-server/Dockerfile index d938ca9936f..f348d293368 100644 --- a/go/infra/docker/feature-server/Dockerfile +++ b/go/infra/docker/feature-server/Dockerfile @@ -1,4 +1,4 @@ -FROM golang:1.26 +FROM golang:1.26-trixie AS builder ENV GOTOOLCHAIN=auto # Update the package list and install the ca-certificates package @@ -24,6 +24,21 @@ RUN find ./protos -name "*.proto" \ # Build the Go application RUN go build -o feast ./go/main.go +# Ship only the server binary on a slim base, so the Go toolchain, compilers and +# kernel headers of the build image are not part of the published image. The +# runtime uses the same Debian release as the builder, so the cgo-linked binary +# finds a matching glibc. +FROM debian:trixie-slim + +# Update the package list and install the ca-certificates package +RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates \ + && rm -rf /var/lib/apt/lists/* + +# Set the current working directory inside the container +WORKDIR /app + +COPY --from=builder /app/feast ./feast + # Expose ports EXPOSE 8080 diff --git a/infra/charts/feast-feature-server/opentelemetry.md b/infra/charts/feast-feature-server/opentelemetry.md index fc0930821e6..cc62859782b 100644 --- a/infra/charts/feast-feature-server/opentelemetry.md +++ b/infra/charts/feast-feature-server/opentelemetry.md @@ -8,7 +8,7 @@ Feast instrumentation Using OpenTelemetry and Prometheus - ### Deploy Prometheus Operator Follow the Prometheus Operator documentation to install the operator - -https://github.com/prometheus-operator/prometheus-operator/blob/main/Documentation/user-guides/getting-started.md +https://github.com/prometheus-operator/prometheus-operator/blob/main/Documentation/getting-started/installation.md ### Deploy OpenTelemetry Operator Before installing OTEL Operator, install `cert-manager` and validate the `pods` should spin up -- diff --git a/protos/feast/core/DataSource.proto b/protos/feast/core/DataSource.proto index 8ea19cebf24..7b1f50fe193 100644 --- a/protos/feast/core/DataSource.proto +++ b/protos/feast/core/DataSource.proto @@ -299,6 +299,9 @@ message DataSource { // the online store on-demand, such as by stream consumers. message PushOptions { reserved 1; + // Names of upstream FeatureViews/LabelViews consumed by the external process feeding this PushSource. + // Represents external pipeline lineage metadata; does not trigger Feast-managed DAG execution. + repeated string upstream_feature_view_names = 2; } diff --git a/sdk/python/feast/api/registry/rest/feature_views.py b/sdk/python/feast/api/registry/rest/feature_views.py index ede3c5d804a..61df2837b4f 100644 --- a/sdk/python/feast/api/registry/rest/feature_views.py +++ b/sdk/python/feast/api/registry/rest/feature_views.py @@ -51,18 +51,25 @@ class ApplyFeatureViewRequestBody(BaseModel): owner: Optional[str] = "" -def _extract_feature_view_from_any(any_feature_view: dict) -> dict: +def _extract_feature_view_from_any( + any_feature_view: dict, project: Optional[str] = None +) -> dict: """Extract the specific feature view type and data from an AnyFeatureView object. Args: any_feature_view: Dictionary containing the AnyFeatureView data + project: Project the feature view was read from, filled into an unset spec.project Returns: Dictionary with 'type' and feature view data, or empty dict if no valid type found """ for key, value in any_feature_view.items(): if value: - return {"type": key, **value} + result = {"type": key, **value} + spec = result.get("spec") + if project and isinstance(spec, dict) and not spec.get("project"): + result["spec"] = {**spec, "project": project} + return result return {} @@ -110,7 +117,9 @@ def list_feature_views_all( response = grpc_call(grpc_handler.ListAllFeatureViews, req) any_feature_views = response.get("featureViews", []) for any_feature_view in any_feature_views: - feature_view = _extract_feature_view_from_any(any_feature_view) + feature_view = _extract_feature_view_from_any( + any_feature_view, project_name + ) if feature_view: feature_view["project"] = project_name all_feature_views.append(feature_view) @@ -157,7 +166,7 @@ def get_any_feature_view( response = grpc_call(grpc_handler.GetAnyFeatureView, req) any_feature_view = response.get("anyFeatureView", {}) - result = _extract_feature_view_from_any(any_feature_view) + result = _extract_feature_view_from_any(any_feature_view, project) if include_relationships: relationships = get_object_relationships( @@ -315,7 +324,7 @@ def list_all_feature_views( # Extract the specific type of feature view from each AnyFeatureView feature_views = [] for any_feature_view in any_feature_views: - feature_view = _extract_feature_view_from_any(any_feature_view) + feature_view = _extract_feature_view_from_any(any_feature_view, project) if feature_view: feature_views.append(feature_view) diff --git a/sdk/python/feast/api/registry/rest/rest_registry_server.py b/sdk/python/feast/api/registry/rest/rest_registry_server.py index d73a7f14672..92122ba46df 100644 --- a/sdk/python/feast/api/registry/rest/rest_registry_server.py +++ b/sdk/python/feast/api/registry/rest/rest_registry_server.py @@ -2,7 +2,7 @@ import logging import re -from fastapi import Depends, FastAPI, HTTPException, Request, status +from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.exceptions import RequestValidationError from fastapi.responses import JSONResponse from pydantic import ValidationError @@ -63,12 +63,6 @@ def __init__(self, store: FeatureStore): root_path="/api/v1", docs_url="/", redoc_url="/docs", - default_status_code=status.HTTP_200_OK, - default_headers={ - "X-Content-Type-Options": "nosniff", - "X-XSS-Protection": "1; mode=block", - "X-Frame-Options": "DENY", - }, ) self._add_exception_handlers() self._add_logging_middleware() diff --git a/sdk/python/feast/data_source.py b/sdk/python/feast/data_source.py index 7d5fca5c1e8..e2f1ddc7b50 100644 --- a/sdk/python/feast/data_source.py +++ b/sdk/python/feast/data_source.py @@ -15,7 +15,7 @@ import warnings from abc import ABC, abstractmethod from datetime import datetime, timedelta, timezone -from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union from google.protobuf.duration_pb2 import Duration from google.protobuf.json_format import MessageToJson @@ -876,12 +876,14 @@ class PushSource(DataSource): # TODO(adchia): consider adding schema here in case where Feast manages pushing events to the offline store # TODO(adchia): consider a "mode" to support pushing raw vs transformed events batch_source: Optional[DataSource] = None + upstream_feature_view_names: List[str] = [] def __init__( self, *, name: str, batch_source: Optional[DataSource] = None, + upstream_feature_view_names: Optional[Sequence[Union[Any, str]]] = None, description: Optional[str] = "", tags: Optional[Dict[str, str]] = None, owner: Optional[str] = "", @@ -890,9 +892,14 @@ def __init__( Creates a PushSource object. Args: - name: Name of the push source - batch_source: The batch source that backs this push source. It's used when materializing from the offline - store to the online store, and when retrieving historical features. + name: Name of the push source. + batch_source (optional): The batch source that backs this push source. It's used when + materializing from the offline store to the online store, and when retrieving + historical features. + upstream_feature_view_names (optional): Sequence of upstream FeatureView/LabelView + objects or names consumed by the external process feeding this push source. + Represents external pipeline lineage metadata and does not trigger Feast-managed + DAG execution. description (optional): A human-readable description. tags (optional): A dictionary of key-value pairs to store arbitrary metadata. owner (optional): The owner of the data source, typically the email of the primary @@ -901,6 +908,27 @@ def __init__( super().__init__(name=name, description=description, tags=tags, owner=owner) self.batch_source = batch_source + normalized_upstream_names: List[str] = [] + if upstream_feature_view_names is not None: + for item in upstream_feature_view_names: + if hasattr(item, "name"): + name_val = item.name + if not isinstance(name_val, str) or not name_val.strip(): + raise ValueError( + f"Upstream feature view object {item} has an invalid or empty name" + ) + normalized_upstream_names.append(name_val.strip()) + elif isinstance(item, str): + if not item.strip(): + raise ValueError("Upstream feature view name cannot be empty") + normalized_upstream_names.append(item.strip()) + else: + raise TypeError( + f"Unsupported type for upstream feature view: {type(item)}. " + "Expected str or object with a 'name' attribute (e.g. FeatureView, LabelView)." + ) + self.upstream_feature_view_names = normalized_upstream_names + def __eq__(self, other): if not isinstance(other, PushSource): return False @@ -911,6 +939,9 @@ def __eq__(self, other): if self.batch_source != other.batch_source: return False + if self.upstream_feature_view_names != other.upstream_feature_view_names: + return False + return True def __hash__(self): @@ -932,9 +963,16 @@ def from_proto(data_source: DataSourceProto): else None ) + upstream_feature_view_names = ( + list(data_source.push_options.upstream_feature_view_names) + if data_source.HasField("push_options") + else [] + ) + return PushSource( name=data_source.name, batch_source=batch_source, + upstream_feature_view_names=upstream_feature_view_names, description=data_source.description, tags=dict(data_source.tags), owner=data_source.owner, @@ -949,6 +987,10 @@ def _to_proto_impl(self) -> DataSourceProto: owner=self.owner, ) + data_source_proto.push_options.upstream_feature_view_names.extend( + self.upstream_feature_view_names + ) + # Only set timestamp fields if we have a batch source and this PushSource doesn't have its own fields if self.batch_source and not ( self.timestamp_field or self.created_timestamp_column or self.field_mapping diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py index ee0e766c11d..917454d19ad 100644 --- a/sdk/python/feast/feature_store.py +++ b/sdk/python/feast/feature_store.py @@ -1474,7 +1474,16 @@ def plan( desired_repo_contents.stream_feature_views, desired_repo_contents.label_views, ) - _validate_data_sources(desired_repo_contents.data_sources) + all_fvs_for_validation = [ + *desired_repo_contents.feature_views, + *desired_repo_contents.on_demand_feature_views, + *desired_repo_contents.stream_feature_views, + *desired_repo_contents.label_views, + ] + _validate_data_sources( + desired_repo_contents.data_sources, + all_feature_views=all_fvs_for_validation, + ) self._make_inferences( desired_repo_contents.data_sources, desired_repo_contents.entities, @@ -1749,7 +1758,7 @@ def apply( # Handle all entityless feature views by using DUMMY_ENTITY as a placeholder entity. entities_to_update.append(DUMMY_ENTITY) - # Validate all feature views and make inferences. + # Validate all feature views and data sources, then make inferences. if not skip_feature_view_validation: self._validate_all_feature_views( views_to_update, @@ -1757,6 +1766,16 @@ def apply( sfvs_to_update, lvs_to_update, ) + all_fvs_for_validation = [ + *views_to_update, + *odfvs_to_update, + *sfvs_to_update, + *(lvs_to_update or []), + ] + _validate_data_sources( + data_sources_to_update, + all_feature_views=all_fvs_for_validation, + ) self._make_inferences( data_sources_to_update, entities_to_update, @@ -5156,8 +5175,21 @@ def _validate_feature_views(feature_views: List[BaseFeatureView]): fv_by_name[case_insensitive_fv_name] = fv -def _validate_data_sources(data_sources: List[DataSource]): - """Verify data sources have case-insensitively unique names.""" +def _validate_data_sources( + data_sources: List[DataSource], + all_feature_views: Optional[ + Sequence[ + Union[ + BaseFeatureView, + FeatureView, + StreamFeatureView, + OnDemandFeatureView, + LabelView, + ] + ] + ] = None, +): + """Verify data sources have case-insensitively unique names and validate PushSource upstream references.""" ds_names = set() for ds in data_sources: case_insensitive_ds_name = ds.name.lower() @@ -5165,3 +5197,16 @@ def _validate_data_sources(data_sources: List[DataSource]): raise DataSourceRepeatNamesException(case_insensitive_ds_name) else: ds_names.add(case_insensitive_ds_name) + + if all_feature_views is not None: + known_fv_names = {fv.name for fv in all_feature_views if hasattr(fv, "name")} + for ds in data_sources: + if isinstance(ds, PushSource) and ds.upstream_feature_view_names: + for upstream_name in ds.upstream_feature_view_names: + if upstream_name not in known_fv_names: + warnings.warn( + f"PushSource '{ds.name}' references upstream feature view '{upstream_name}' " + "which was not found in the current feature repository.", + UserWarning, + stacklevel=2, + ) diff --git a/sdk/python/feast/infra/compute_engines/ray/feature_builder.py b/sdk/python/feast/infra/compute_engines/ray/feature_builder.py index f0a26df263b..d88354c2fc2 100644 --- a/sdk/python/feast/infra/compute_engines/ray/feature_builder.py +++ b/sdk/python/feast/infra/compute_engines/ray/feature_builder.py @@ -88,6 +88,9 @@ def build_aggregation_node(self, view, input_node: DAGNode) -> DAGNode: config=self.config, enable_tiling=enable_tiling, hop_size=tiling_hop_size, + created_timestamp_col=getattr( + view.batch_source, "created_timestamp_column", None + ), ) node.add_input(input_node) @@ -333,19 +336,19 @@ def _build_aggregation_optimized_plan(self) -> ExecutionPlan: # 1. Read source data last_node = self.build_source_node(self.feature_view) - # 2. Apply filters (TTL, custom filters) BEFORE aggregation + # 2. Apply transformations to the raw source data, as materialization does + if self._should_transform(self.feature_view): + last_node = self.build_transformation_node(self.feature_view, [last_node]) + + # 3. Apply filters (TTL, custom filters) BEFORE aggregation last_node = self.build_filter_node(self.feature_view, last_node) - # 3. Aggregate across all historical records + # 4. Aggregate across all historical records last_node = self.build_aggregation_node(self.feature_view, last_node) - # 4. Join with entity_df to get aggregated features for each entity + # 5. Join with entity_df to get aggregated features for each entity last_node = self.build_join_node(self.feature_view, [last_node]) - # 5. Apply transformations to aggregated features - if self._should_transform(self.feature_view): - last_node = self.build_transformation_node(self.feature_view, [last_node]) - # 6. Output last_node = self.build_output_nodes(self.feature_view, last_node) diff --git a/sdk/python/feast/infra/compute_engines/ray/nodes.py b/sdk/python/feast/infra/compute_engines/ray/nodes.py index edc6dafd5cd..1b421a8a3d2 100644 --- a/sdk/python/feast/infra/compute_engines/ray/nodes.py +++ b/sdk/python/feast/infra/compute_engines/ray/nodes.py @@ -20,6 +20,7 @@ from feast.infra.compute_engines.dag.value import DAGValue from feast.infra.compute_engines.ray.config import RayComputeEngineConfig from feast.infra.compute_engines.ray.utils import ( + _is_empty_batch, safe_batch_processor, write_to_online_store, ) @@ -329,11 +330,13 @@ def __init__( config: RayComputeEngineConfig, enable_tiling: bool = False, hop_size: Optional[timedelta] = None, + created_timestamp_col: Optional[str] = None, ): super().__init__(name) self.aggregations = aggregations self.group_by_keys = group_by_keys self.timestamp_col = timestamp_col + self.created_timestamp_col = created_timestamp_col self.config = config self.enable_tiling = enable_tiling self.hop_size = hop_size @@ -443,6 +446,21 @@ def _execute_tiled_aggregation(self, dataset: Dataset) -> DAGValue: final_pdf, self.group_by_keys, self.timestamp_col ) + # Tiling drops the created timestamp; add back the latest per group. + if ( + self.created_timestamp_col + and self.created_timestamp_col in input_pdf.columns + and not final_pdf.empty + ): + latest_created = ( + input_pdf.groupby(self.group_by_keys)[self.created_timestamp_col] + .max() + .reset_index() + ) + final_pdf = final_pdf.merge( + latest_created, on=self.group_by_keys, how="left" + ) + aggregated_dataset = ray_wrapper.from_pandas(final_pdf) return DAGValue( @@ -543,6 +561,10 @@ def _fallback_pandas_aggregation(self, dataset: Dataset, agg_dict: dict) -> Data # Combine aggregation results if agg_results: + # Keep the latest timestamps per group; the online write needs them. + for ts_col in (self.timestamp_col, self.created_timestamp_col): + if ts_col and ts_col in df.columns: + agg_results.append(grouped[ts_col].max()) result_df = pd.concat(agg_results, axis=1) # Reset index to make group keys regular columns @@ -708,11 +730,14 @@ def execute(self, context: ExecutionContext) -> DAGValue: elif callable(self.transformation): transformation_serialized = dill.dumps(self.transformation) - @safe_batch_processor + # Not wrapped in safe_batch_processor: a failed UDF must fail the job. def apply_transformation_with_serialized_udf( batch: pd.DataFrame, ) -> pd.DataFrame: """Apply the transformation using pre-serialized UDF.""" + if _is_empty_batch(batch): + return batch + if transformation_serialized: transformation_func = dill.loads(transformation_serialized) transformed_batch = transformation_func(batch) @@ -861,9 +886,12 @@ def execute(self, context: ExecutionContext) -> DAGValue: feature_view=self.feature_view, repo_config=context.repo_config ) - @safe_batch_processor + # Not wrapped in safe_batch_processor: a failed write must fail the job. def write_batch_with_serialized_artifacts(batch: pd.DataFrame) -> pd.DataFrame: """Write each batch using pre-serialized artifacts.""" + if _is_empty_batch(batch): + return batch + ( feature_view, online_store, diff --git a/sdk/python/feast/infra/compute_engines/ray/utils.py b/sdk/python/feast/infra/compute_engines/ray/utils.py index 9727f010ca4..2a4f85e2274 100644 --- a/sdk/python/feast/infra/compute_engines/ray/utils.py +++ b/sdk/python/feast/infra/compute_engines/ray/utils.py @@ -81,6 +81,7 @@ def write_to_online_store( except Exception as e: logger.error(f"Failed to write to online store for {feature_view.name}: {e}") + raise # Ray Data batch type: pandas DataFrame, numpy dict, or pyarrow Table @@ -178,6 +179,7 @@ def _write_batch(batch: pa.Table) -> pa.Table: logger.error( f"Distributed write to online store failed for {feature_view.name}: {e}" ) + raise def safe_batch_processor( diff --git a/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/__init__.py b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/__init__.py index 2a6d36f87e2..5c3856ecb00 100644 --- a/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/__init__.py +++ b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/__init__.py @@ -9,6 +9,12 @@ from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import ( IcebergSource, ) +from feast.infra.data_sources.contrib.iceberg_catalog.polaris_client import ( + PolarisClient, +) +from feast.infra.data_sources.contrib.iceberg_catalog.polaris_source import ( + PolarisSource, +) from feast.infra.data_sources.contrib.iceberg_catalog.uc_client import UCClient from feast.infra.data_sources.contrib.iceberg_catalog.unity_catalog_source import ( UnityCatalogSource, @@ -19,6 +25,8 @@ "IcebergCatalogError", "IcebergRestClient", "IcebergSource", + "PolarisClient", + "PolarisSource", "TableMetadata", "TableNotFoundError", "TempCredential", diff --git a/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/iceberg_source.py b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/iceberg_source.py index c93b5db62d1..94a9a38a965 100644 --- a/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/iceberg_source.py +++ b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/iceberg_source.py @@ -131,6 +131,38 @@ def get_catalog_client(self): def source_type(self) -> DataSourceProto.SourceType.ValueType: return DataSourceProto.BATCH_ICEBERG + def get_pyiceberg_catalog_config(self) -> Dict[str, str]: + """Return the PyIceberg ``load_catalog`` properties for a REST catalog. + + Subclasses (for example ``PolarisSource``) override this to add + catalog-specific authentication or credential-vending properties. + """ + config: Dict[str, str] = { + "type": "rest", + "uri": self.endpoint, + "warehouse": self.warehouse, + } + if self.token_env_var: + token = os.environ.get(self.token_env_var) + if token: + config["token"] = token + return config + + def load_pyiceberg_table(self): + """Load the backing table as a PyIceberg ``Table``. + + REST sources are loaded from ``get_pyiceberg_catalog_config()``; other + catalog types delegate to ``get_catalog_client()``. + """ + fqn = f"{self.namespace}.{self.iceberg_table}" + if self.catalog_type != "rest": + return self.get_catalog_client().load_table(fqn) + + from pyiceberg.catalog import load_catalog + + catalog = load_catalog(self.catalog_name, **self.get_pyiceberg_catalog_config()) + return catalog.load_table(fqn) + @staticmethod def from_proto(data_source: DataSourceProto) -> Any: cls_type = data_source.data_source_class_type @@ -140,6 +172,12 @@ def from_proto(data_source: DataSourceProto) -> Any: ) return UnityCatalogSource.from_proto(data_source) + if cls_type and "PolarisSource" in cls_type: + from feast.infra.data_sources.contrib.iceberg_catalog.polaris_source import ( + PolarisSource, + ) + + return PolarisSource.from_proto(data_source) custom_options = json.loads( str(data_source.custom_options.configuration, encoding="utf8") diff --git a/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_client.py b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_client.py new file mode 100644 index 00000000000..f2800c9c085 --- /dev/null +++ b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_client.py @@ -0,0 +1,216 @@ +# Copyright 2026 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import threading +import time +from typing import Dict, List, Optional + +import requests +from requests.auth import AuthBase + +from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_rest_client import ( + CatalogAuthError, + IcebergRestClient, +) + +logger = logging.getLogger(__name__) + +DEFAULT_OAUTH_SCOPE = "PRINCIPAL_ROLE:ALL" +_DEFAULT_TOKEN_TTL_SECONDS = 3600 + + +class OAuthClientCredentialsTokenProvider: + """Fetches and caches OAuth2 client-credentials access tokens. + + The token is re-fetched shortly before it expires so that long-running + jobs (materialization, feature serving) keep working. + """ + + def __init__( + self, + token_url: str, + client_id: str, + client_secret: str, + scope: str = DEFAULT_OAUTH_SCOPE, + timeout: int = 30, + refresh_skew_seconds: int = 60, + ): + self.token_url = token_url + self._client_id = client_id + self._client_secret = client_secret + self.scope = scope + self.timeout = timeout + self.refresh_skew_seconds = refresh_skew_seconds + self._token: Optional[str] = None + self._expires_at: float = 0.0 + self._lock = threading.Lock() + + def get_token(self, force_refresh: bool = False) -> str: + with self._lock: + now = time.monotonic() + if ( + force_refresh + or self._token is None + or now >= self._expires_at - self.refresh_skew_seconds + ): + self._token, ttl = self._fetch_token() + self._expires_at = now + ttl + return self._token + + def _fetch_token(self) -> "tuple[str, float]": + try: + response = requests.post( + self.token_url, + data={ + "grant_type": "client_credentials", + "client_id": self._client_id, + "client_secret": self._client_secret, + "scope": self.scope, + }, + headers={"Accept": "application/json"}, + timeout=self.timeout, + ) + except requests.exceptions.RequestException as e: + raise CatalogAuthError( + f"Failed to reach OAuth token endpoint {self.token_url}: {e}" + ) from e + + if not response.ok: + raise CatalogAuthError( + f"OAuth token request to {self.token_url} failed: " + f"HTTP {response.status_code}", + status_code=response.status_code, + ) + body = response.json() + token = body.get("access_token") + if not token: + raise CatalogAuthError( + f"OAuth token response from {self.token_url} has no access_token" + ) + return token, float(body.get("expires_in", _DEFAULT_TOKEN_TTL_SECONDS)) + + +class _BearerTokenAuth(AuthBase): + """requests auth hook that attaches a fresh bearer token to every request.""" + + def __init__(self, provider: OAuthClientCredentialsTokenProvider): + self._provider = provider + + def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest: + request.headers["Authorization"] = f"Bearer {self._provider.get_token()}" + return request + + +class PolarisClient(IcebergRestClient): + """Apache Polaris client. + + Extends :class:`IcebergRestClient` with: + + - OAuth2 client-credentials authentication with automatic token refresh + (``client_id`` / ``client_secret``), as an alternative to a static token. + - Table property management through the standard Iceberg REST + ``updateTable`` operation (``set-properties`` / ``remove-properties``), + used to record Feast metadata on the Polaris table. + + Polaris serves the Iceberg REST API under ``/api/catalog`` and uses + the Polaris catalog name as the ``/v1/{prefix}`` path segment. + """ + + def __init__( + self, + endpoint: str, + warehouse: str, + token: Optional[str] = None, + client_id: Optional[str] = None, + client_secret: Optional[str] = None, + scope: str = DEFAULT_OAUTH_SCOPE, + oauth_server_uri: Optional[str] = None, + credential_vending: bool = True, + timeout: int = 30, + ): + super().__init__( + endpoint=endpoint, + token=None if (client_id and client_secret) else token, + warehouse=warehouse, + credential_vending=credential_vending, + timeout=timeout, + url_style="prefix", + ) + self.token_provider: Optional[OAuthClientCredentialsTokenProvider] = None + if client_id and client_secret: + self.token_provider = OAuthClientCredentialsTokenProvider( + token_url=oauth_server_uri or f"{self.endpoint}/v1/oauth/tokens", + client_id=client_id, + client_secret=client_secret, + scope=scope, + timeout=timeout, + ) + self._session.auth = _BearerTokenAuth(self.token_provider) + + def set_table_properties( + self, namespace: str, table: str, properties: Dict[str, str] + ) -> bool: + """Set properties on a table. Returns True on success. + + Failures are logged and reported as ``False`` so that metadata sync + never aborts ``feast apply``. + """ + if not properties: + return True + return self._commit_updates( + namespace, + table, + [ + { + "action": "set-properties", + "updates": {k: str(v) for k, v in properties.items()}, + } + ], + ) + + def remove_table_properties( + self, namespace: str, table: str, keys: List[str] + ) -> bool: + """Remove properties from a table. Returns True on success.""" + if not keys: + return True + return self._commit_updates( + namespace, table, [{"action": "remove-properties", "removals": keys}] + ) + + def _commit_updates( + self, namespace: str, table: str, updates: List[Dict[str, object]] + ) -> bool: + url = f"{self._base_url}/namespaces/{namespace}/tables/{table}" + try: + response = self._session.post( + url, + json={"requirements": [], "updates": updates}, + timeout=self.timeout, + ) + except (requests.exceptions.RequestException, CatalogAuthError) as e: + logger.warning("Failed to update table %s.%s: %s", namespace, table, e) + return False + + if not response.ok: + logger.warning( + "Failed to update table %s.%s: HTTP %s %s", + namespace, + table, + response.status_code, + response.text[:200], + ) + return False + return True diff --git a/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_provider.py b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_provider.py new file mode 100644 index 00000000000..0891b193d24 --- /dev/null +++ b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_provider.py @@ -0,0 +1,125 @@ +# Copyright 2026 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Polaris Provider: records Feast metadata on Polaris tables. + +Detects :class:`PolarisSource` in feature views during ``feast apply`` and sets +``feast.*`` table properties through the standard Iceberg REST ``updateTable`` +operation. Tables are never dropped; when a feature view is deleted the +properties are updated to mark the table as no longer a feature table. + +Usage in feature_store.yaml:: + + provider: polaris +""" + +import logging +from datetime import datetime, timezone +from typing import Dict, Sequence + +from feast.base_feature_view import BaseFeatureView +from feast.entity import Entity +from feast.feature_view import FeatureView +from feast.infra.data_sources.contrib.iceberg_catalog.polaris_source import ( + PolarisSource, +) +from feast.infra.passthrough_provider import PassthroughProvider + +logger = logging.getLogger(__name__) + +TAG_PREFIX = "feast.tag." + + +class PolarisProvider(PassthroughProvider): + """Provider that syncs Feast metadata to Polaris on ``feast apply``.""" + + def update_infra( + self, + project: str, + tables_to_delete: Sequence[BaseFeatureView], + tables_to_keep: Sequence[BaseFeatureView], + entities_to_delete: Sequence[Entity], + entities_to_keep: Sequence[Entity], + partial: bool, + ): + super().update_infra( + project, + tables_to_delete, + tables_to_keep, + entities_to_delete, + entities_to_keep, + partial, + ) + self._sync_polaris_properties(project, tables_to_keep) + self._mark_polaris_tables_removed(tables_to_delete) + + @staticmethod + def _polaris_source(view: BaseFeatureView): + if not isinstance(view, FeatureView): + return None + source = view.batch_source + if isinstance(source, PolarisSource) and source.sync_metadata: + return source + return None + + @staticmethod + def build_properties(project: str, view: FeatureView) -> Dict[str, str]: + properties: Dict[str, str] = { + "feast.is_feature_table": "true", + "feast.project": project, + "feast.feature_view": view.name, + } + if view.owner: + properties["feast.owner"] = view.owner + if view.description: + properties["feast.description"] = view.description + for key, value in (view.tags or {}).items(): + properties[f"{TAG_PREFIX}{key}"] = value + return properties + + def _sync_polaris_properties( + self, project: str, views: Sequence[BaseFeatureView] + ) -> None: + for view in views: + source = self._polaris_source(view) + if source is None: + continue + assert isinstance(view, FeatureView) + logger.info("Syncing Feast metadata for %s to %s", view.name, source.fqn) + client = source.get_catalog_client() + client.set_table_properties( + source.namespace, + source.iceberg_table, + self.build_properties(project, view), + ) + + def _mark_polaris_tables_removed(self, views: Sequence[BaseFeatureView]) -> None: + for view in views: + source = self._polaris_source(view) + if source is None: + continue + logger.info( + "Feature view %s deleted; keeping Polaris table %s and marking it", + view.name, + source.fqn, + ) + client = source.get_catalog_client() + client.set_table_properties( + source.namespace, + source.iceberg_table, + { + "feast.is_feature_table": "false", + "feast.deleted_at": datetime.now(timezone.utc).isoformat(), + }, + ) diff --git a/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_source.py b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_source.py new file mode 100644 index 00000000000..099b4b6a69e --- /dev/null +++ b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/polaris_source.py @@ -0,0 +1,298 @@ +# Copyright 2026 The Feast Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import logging +import os +from typing import Any, Dict, List, Optional, Sequence, Tuple + +from feast.credentials import ConnectionRef +from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import ( + IcebergSource, +) +from feast.infra.data_sources.contrib.iceberg_catalog.polaris_client import ( + DEFAULT_OAUTH_SCOPE, + PolarisClient, +) +from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto + +logger = logging.getLogger(__name__) + +DEFAULT_TIMESTAMP_CANDIDATES = ("event_timestamp", "event_ts", "timestamp", "ts") + + +class PolarisSource(IcebergSource): + """Data source backed by an Apache Polaris catalog. + + Extends :class:`IcebergSource` with Polaris-specific behaviour: + + - ``endpoint`` defaults to ``$POLARIS_HOST/api/catalog``. + - Authentication via OAuth2 client credentials (read from the environment + variables named by ``client_id_env_var`` / ``client_secret_env_var``, + with automatic token refresh), falling back to a static bearer token in + ``token_env_var``. Only environment variable *names* are stored in the + registry, never secrets. + - Credential vending (``X-Iceberg-Access-Delegation``) is requested for + reads and writes done through PyIceberg when ``credential_vending`` is set. + - With ``sync_metadata`` set, ``PolarisProvider`` records Feast metadata + (project, feature view, owner, tags) as Polaris table properties on + ``feast apply``. + + ``warehouse`` is the Polaris catalog name. + """ + + def __init__( + self, + *, + warehouse: str, + namespace: str, + table: str, + name: Optional[str] = None, + timestamp_field: Optional[str] = None, + created_timestamp_column: Optional[str] = None, + field_mapping: Optional[Dict[str, str]] = None, + description: Optional[str] = "", + tags: Optional[Dict[str, str]] = None, + owner: Optional[str] = "", + endpoint: Optional[str] = None, + token_env_var: Optional[str] = None, + client_id_env_var: Optional[str] = None, + client_secret_env_var: Optional[str] = None, + oauth_scope: str = DEFAULT_OAUTH_SCOPE, + oauth_server_uri: Optional[str] = None, + credential_vending: bool = True, + catalog_name: str = "feast_polaris", + sync_metadata: bool = True, + connection_ref: Optional[ConnectionRef] = None, + ): + if endpoint is None: + host = os.environ.get("POLARIS_HOST", "").rstrip("/") + if not host: + raise ValueError( + "PolarisSource needs an endpoint: pass `endpoint=` or set " + "the POLARIS_HOST environment variable." + ) + endpoint = f"{host}/api/catalog" + + super().__init__( + warehouse=warehouse, + namespace=namespace, + table=table, + catalog_type="rest", + catalog_name=catalog_name, + endpoint=endpoint, + name=name, + timestamp_field=timestamp_field, + created_timestamp_column=created_timestamp_column, + field_mapping=field_mapping, + description=description, + tags=tags, + owner=owner, + token_env_var=token_env_var or "POLARIS_TOKEN", + credential_vending=credential_vending, + connection_ref=connection_ref, + ) + self.client_id_env_var = client_id_env_var or "POLARIS_CLIENT_ID" + self.client_secret_env_var = client_secret_env_var or "POLARIS_CLIENT_SECRET" + self.oauth_scope = oauth_scope + self.oauth_server_uri = oauth_server_uri + self.sync_metadata = sync_metadata + + def _resolve_client_credentials(self) -> Tuple[Optional[str], Optional[str]]: + client_id = os.environ.get(self.client_id_env_var) + client_secret = os.environ.get(self.client_secret_env_var) + if client_id and client_secret: + return client_id, client_secret + return None, None + + def get_catalog_client(self) -> PolarisClient: # type: ignore[override] + client_id, client_secret = self._resolve_client_credentials() + token = os.environ.get(self.token_env_var) if self.token_env_var else None + if not (client_id and client_secret) and not token: + logger.warning( + "No Polaris credentials found in the environment for source " + "%s (expected OAuth client credentials or a bearer token); " + "requests will be unauthenticated.", + self.name, + ) + return PolarisClient( + endpoint=self.endpoint, + warehouse=self.warehouse, + token=token, + client_id=client_id, + client_secret=client_secret, + scope=self.oauth_scope, + oauth_server_uri=self.oauth_server_uri, + credential_vending=self.credential_vending, + ) + + def get_pyiceberg_catalog_config(self) -> Dict[str, str]: + config: Dict[str, str] = { + "type": "rest", + "uri": self.endpoint, + "warehouse": self.warehouse, + } + client_id, client_secret = self._resolve_client_credentials() + if client_id and client_secret: + config["credential"] = f"{client_id}:{client_secret}" + config["scope"] = self.oauth_scope + if self.oauth_server_uri: + config["oauth2-server-uri"] = self.oauth_server_uri + elif self.token_env_var: + token = os.environ.get(self.token_env_var) + if token: + config["token"] = token + if self.credential_vending: + config["header.X-Iceberg-Access-Delegation"] = "vended-credentials" + return config + + @staticmethod + def discover( + *, + warehouse: str, + endpoint: Optional[str] = None, + namespaces: Optional[Sequence[str]] = None, + timestamp_fields: Sequence[str] = DEFAULT_TIMESTAMP_CANDIDATES, + **source_kwargs: Any, + ) -> List["PolarisSource"]: + """Build a ``PolarisSource`` for every table in the catalog. + + Tables are skipped unless they have a column named in + ``timestamp_fields`` (the first match is used as ``timestamp_field``). + Extra keyword arguments (``token_env_var``, ``client_id_env_var``, + ``tags``, ...) are passed to every created source. Each source is named + ``__``. + """ + probe = PolarisSource( + warehouse=warehouse, + namespace="_", + table="_", + endpoint=endpoint, + **source_kwargs, + ) + client = probe.get_catalog_client() + sources: List[PolarisSource] = [] + for namespace in namespaces or client.list_namespaces(): + for table in client.list_tables(namespace): + columns = {f["name"] for f in client.get_table_schema(namespace, table)} + timestamp_field = next( + (c for c in timestamp_fields if c in columns), None + ) + if timestamp_field is None: + logger.info( + "Skipping %s.%s: no timestamp column among %s", + namespace, + table, + list(timestamp_fields), + ) + continue + sources.append( + PolarisSource( + name=f"{namespace}__{table}", + warehouse=warehouse, + namespace=namespace, + table=table, + endpoint=probe.endpoint, + timestamp_field=timestamp_field, + **source_kwargs, + ) + ) + return sources + + @staticmethod + def from_proto(data_source: DataSourceProto) -> Any: + options = json.loads( + str(data_source.custom_options.configuration, encoding="utf8") + ) + return PolarisSource( + name=data_source.name, + endpoint=options["endpoint"], + warehouse=options["warehouse"], + namespace=options["namespace"], + table=options["table"], + token_env_var=options.get("token_env_var"), + client_id_env_var=options.get("client_id_env_var"), + client_secret_env_var=options.get("client_secret_env_var"), + oauth_scope=options.get("oauth_scope", DEFAULT_OAUTH_SCOPE), + oauth_server_uri=options.get("oauth_server_uri"), + credential_vending=options.get("credential_vending", True), + catalog_name=options.get("catalog_name", "feast_polaris"), + sync_metadata=options.get("sync_metadata", True), + timestamp_field=data_source.timestamp_field, + created_timestamp_column=data_source.created_timestamp_column, + field_mapping=dict(data_source.field_mapping), + description=data_source.description, + tags=dict(data_source.tags), + owner=data_source.owner, + connection_ref=ConnectionRef.from_tags(dict(data_source.tags)), + ) + + def _to_proto_impl(self) -> DataSourceProto: + config_json = json.dumps( + { + "endpoint": self.endpoint, + "warehouse": self.warehouse, + "namespace": self.namespace, + "table": self.iceberg_table, + "token_env_var": self.token_env_var, + "client_id_env_var": self.client_id_env_var, + "client_secret_env_var": self.client_secret_env_var, + "oauth_scope": self.oauth_scope, + "oauth_server_uri": self.oauth_server_uri, + "credential_vending": self.credential_vending, + "catalog_name": self.catalog_name, + "sync_metadata": self.sync_metadata, + } + ) + return DataSourceProto( + name=self.name, + type=DataSourceProto.BATCH_ICEBERG, + data_source_class_type=( + "feast.infra.data_sources.contrib.iceberg_catalog" + ".polaris_source.PolarisSource" + ), + custom_options=DataSourceProto.CustomSourceOptions( + configuration=bytes(config_json, encoding="utf8") + ), + description=self.description, + tags=self.tags, + owner=self.owner, + timestamp_field=self.timestamp_field, + created_timestamp_column=self.created_timestamp_column, + field_mapping=self.field_mapping, + ) + + def __eq__(self, other): + if not isinstance(other, PolarisSource): + return False + if not super().__eq__(other): + return False + return ( + self.client_id_env_var == other.client_id_env_var + and self.client_secret_env_var == other.client_secret_env_var + and self.oauth_scope == other.oauth_scope + and self.oauth_server_uri == other.oauth_server_uri + and self.sync_metadata == other.sync_metadata + ) + + def __hash__(self): + return hash( + ( + self.name, + self.endpoint, + self.warehouse, + self.namespace, + self.iceberg_table, + ) + ) diff --git a/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/tests/test_polaris.py b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/tests/test_polaris.py new file mode 100644 index 00000000000..a78b256e6f1 --- /dev/null +++ b/sdk/python/feast/infra/data_sources/contrib/iceberg_catalog/tests/test_polaris.py @@ -0,0 +1,430 @@ +"""Unit tests for PolarisSource, PolarisClient and PolarisProvider.""" + +import json +from typing import Any, Dict +from unittest.mock import MagicMock, Mock, patch + +import pytest +import requests + +from feast.entity import Entity +from feast.feature_view import FeatureView +from feast.field import Field +from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_rest_client import ( + CatalogAuthError, +) +from feast.infra.data_sources.contrib.iceberg_catalog.iceberg_source import ( + IcebergSource, +) +from feast.infra.data_sources.contrib.iceberg_catalog.polaris_client import ( + OAuthClientCredentialsTokenProvider, + PolarisClient, +) +from feast.infra.data_sources.contrib.iceberg_catalog.polaris_provider import ( + PolarisProvider, +) +from feast.infra.data_sources.contrib.iceberg_catalog.polaris_source import ( + PolarisSource, +) +from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto +from feast.types import Float32 +from feast.value_type import ValueType + +ENDPOINT = "http://polaris:8181/api/catalog" + + +def _source(**kwargs) -> PolarisSource: + defaults: Dict[str, Any] = dict( + endpoint=ENDPOINT, + warehouse="my_catalog", + namespace="ml", + table="driver_stats", + timestamp_field="event_timestamp", + ) + defaults.update(kwargs) + return PolarisSource(**defaults) + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + for var in ( + "POLARIS_HOST", + "POLARIS_TOKEN", + "POLARIS_CLIENT_ID", + "POLARIS_CLIENT_SECRET", + ): + monkeypatch.delenv(var, raising=False) + + +class TestPolarisSource: + def test_defaults(self): + source = _source() + assert source.catalog_type == "rest" + assert source.token_env_var == "POLARIS_TOKEN" + assert source.client_id_env_var == "POLARIS_CLIENT_ID" + secret_var = source.client_secret_env_var + assert secret_var == "POLARIS_CLIENT_SECRET" # pragma: allowlist secret + assert source.sync_metadata is True + assert source.credential_vending is True + assert source.fqn == "my_catalog.ml.driver_stats" + + def test_is_subclass_of_iceberg_source(self): + assert isinstance(_source(), IcebergSource) + + def test_endpoint_from_polaris_host(self, monkeypatch): + monkeypatch.setenv("POLARIS_HOST", "http://polaris:8181/") + source = PolarisSource(warehouse="w", namespace="n", table="t") + assert source.endpoint == "http://polaris:8181/api/catalog" + + def test_missing_endpoint_raises(self): + with pytest.raises(ValueError, match="POLARIS_HOST"): + PolarisSource(warehouse="w", namespace="n", table="t") + + def test_proto_roundtrip(self): + source = _source( + description="d", + tags={"team": "ml"}, + owner="o@example.com", + client_id_env_var="MY_ID", + client_secret_env_var="MY_SECRET", # pragma: allowlist secret + oauth_scope="PRINCIPAL_ROLE:feast", + oauth_server_uri="http://idp/token", + credential_vending=False, + sync_metadata=False, + ) + proto = source.to_proto() + assert proto.type == DataSourceProto.BATCH_ICEBERG + assert proto.data_source_class_type.endswith("polaris_source.PolarisSource") + options = json.loads(str(proto.custom_options.configuration, "utf8")) + # Only env var names are persisted, never secrets. + assert options["client_id_env_var"] == "MY_ID" + assert "client_secret" not in options + + restored = PolarisSource.from_proto(proto) + assert isinstance(restored, PolarisSource) + assert restored == source + assert restored.oauth_scope == "PRINCIPAL_ROLE:feast" + assert restored.oauth_server_uri == "http://idp/token" + assert restored.sync_metadata is False + assert restored.credential_vending is False + + def test_iceberg_source_from_proto_dispatches_to_polaris(self): + proto = _source().to_proto() + restored = IcebergSource.from_proto(proto) + assert isinstance(restored, PolarisSource) + + def test_equality_considers_polaris_fields(self): + assert _source() == _source() + assert _source(sync_metadata=False) != _source() + assert _source(oauth_scope="PRINCIPAL_ROLE:x") != _source() + + def test_pyiceberg_config_static_token(self, monkeypatch): + monkeypatch.setenv("POLARIS_TOKEN", "tok") + config = _source().get_pyiceberg_catalog_config() + assert config["type"] == "rest" + assert config["uri"] == ENDPOINT + assert config["warehouse"] == "my_catalog" + assert config["token"] == "tok" + assert "credential" not in config + assert config["header.X-Iceberg-Access-Delegation"] == "vended-credentials" + + def test_pyiceberg_config_oauth_wins_over_token(self, monkeypatch): + monkeypatch.setenv("POLARIS_TOKEN", "tok") + monkeypatch.setenv("POLARIS_CLIENT_ID", "cid") + monkeypatch.setenv("POLARIS_CLIENT_SECRET", "sec") + config = _source( + oauth_server_uri="http://idp/token" + ).get_pyiceberg_catalog_config() + assert config["credential"] == "cid:sec" + assert config["scope"] == "PRINCIPAL_ROLE:ALL" + assert config["oauth2-server-uri"] == "http://idp/token" + assert "token" not in config + + def test_pyiceberg_config_without_credential_vending(self, monkeypatch): + monkeypatch.setenv("POLARIS_TOKEN", "tok") + config = _source(credential_vending=False).get_pyiceberg_catalog_config() + assert "header.X-Iceberg-Access-Delegation" not in config + + @patch("pyiceberg.catalog.load_catalog") + def test_load_pyiceberg_table_uses_polaris_config(self, mock_load, monkeypatch): + monkeypatch.setenv("POLARIS_CLIENT_ID", "cid") + monkeypatch.setenv("POLARIS_CLIENT_SECRET", "sec") + _source().load_pyiceberg_table() + name, kwargs = mock_load.call_args[0][0], mock_load.call_args[1] + assert name == "feast_polaris" + assert kwargs["credential"] == "cid:sec" + mock_load.return_value.load_table.assert_called_once_with("ml.driver_stats") + + def test_catalog_client_is_polaris_client(self, monkeypatch): + monkeypatch.setenv("POLARIS_CLIENT_ID", "cid") + monkeypatch.setenv("POLARIS_CLIENT_SECRET", "sec") + client = _source().get_catalog_client() + assert isinstance(client, PolarisClient) + assert client.token_provider is not None + + def test_base_iceberg_source_config_is_unchanged(self, monkeypatch): + monkeypatch.setenv("T", "tok") + source = IcebergSource( + endpoint="http://host/iceberg", + warehouse="w", + namespace="n", + table="t", + token_env_var="T", + ) + assert source.get_pyiceberg_catalog_config() == { + "type": "rest", + "uri": "http://host/iceberg", + "warehouse": "w", + "token": "tok", + } + + +class TestDiscover: + @patch.object(PolarisClient, "get_table_schema") + @patch.object(PolarisClient, "list_tables") + @patch.object(PolarisClient, "list_namespaces") + def test_discover_builds_sources_and_skips_tables_without_timestamp( + self, mock_ns, mock_tables, mock_schema + ): + mock_ns.return_value = ["ml"] + mock_tables.return_value = ["with_ts", "without_ts"] + mock_schema.side_effect = lambda ns, t: ( + [ + {"name": "id", "type": "long"}, + {"name": "event_timestamp", "type": "timestamptz"}, + ] + if t == "with_ts" + else [{"name": "id", "type": "long"}] + ) + sources = PolarisSource.discover( + warehouse="my_catalog", endpoint=ENDPOINT, tags={"team": "ml"} + ) + assert [s.name for s in sources] == ["ml__with_ts"] + assert sources[0].timestamp_field == "event_timestamp" + assert sources[0].tags == {"team": "ml"} + assert sources[0].endpoint == ENDPOINT + + @patch.object(PolarisClient, "list_tables") + @patch.object(PolarisClient, "list_namespaces") + def test_discover_respects_explicit_namespaces(self, mock_ns, mock_tables): + mock_tables.return_value = [] + PolarisSource.discover(warehouse="w", endpoint=ENDPOINT, namespaces=["a", "b"]) + mock_ns.assert_not_called() + assert [c.args[0] for c in mock_tables.call_args_list] == ["a", "b"] + + +class TestPolarisClientAuth: + def test_static_token_header(self): + client = PolarisClient(ENDPOINT, "w", token="tok") + assert client._session.headers["Authorization"] == "Bearer tok" + assert client.token_provider is None + assert client.url_style == "prefix" + + def test_oauth_overrides_static_token(self): + client = PolarisClient( + ENDPOINT, "w", token="tok", client_id="cid", client_secret="sec" + ) + assert "Authorization" not in client._session.headers + assert client.token_provider is not None + assert client.token_provider.token_url == f"{ENDPOINT}/v1/oauth/tokens" + + @patch("requests.post") + def test_token_fetch_and_cache(self, mock_post): + mock_post.return_value = Mock( + ok=True, json=lambda: {"access_token": "abc", "expires_in": 3600} + ) + provider = OAuthClientCredentialsTokenProvider( + "http://idp/token", "cid", "sec", scope="PRINCIPAL_ROLE:ALL" + ) + assert provider.get_token() == "abc" + assert provider.get_token() == "abc" + assert mock_post.call_count == 1 + data = mock_post.call_args.kwargs["data"] + assert data["grant_type"] == "client_credentials" + assert data["client_id"] == "cid" + assert data["scope"] == "PRINCIPAL_ROLE:ALL" + + @patch("requests.post") + def test_token_refreshes_before_expiry(self, mock_post): + mock_post.side_effect = [ + Mock(ok=True, json=lambda: {"access_token": "t1", "expires_in": 100}), + Mock(ok=True, json=lambda: {"access_token": "t2", "expires_in": 100}), + ] + provider = OAuthClientCredentialsTokenProvider( + "http://idp/token", "cid", "sec", refresh_skew_seconds=60 + ) + with patch( + "feast.infra.data_sources.contrib.iceberg_catalog.polaris_client.time.monotonic" + ) as clock: + clock.return_value = 1000.0 + assert provider.get_token() == "t1" + clock.return_value = 1030.0 # 70s left, outside the 60s skew + assert provider.get_token() == "t1" + clock.return_value = 1050.0 # 50s left, inside the skew -> refresh + assert provider.get_token() == "t2" + assert mock_post.call_count == 2 + + @patch("requests.post") + def test_token_http_error_raises_auth_error(self, mock_post): + mock_post.return_value = Mock(ok=False, status_code=401) + provider = OAuthClientCredentialsTokenProvider("http://idp/token", "c", "s") + with pytest.raises(CatalogAuthError): + provider.get_token() + + @patch("requests.post") + def test_token_connection_error_raises_auth_error(self, mock_post): + mock_post.side_effect = requests.exceptions.ConnectionError("down") + provider = OAuthClientCredentialsTokenProvider("http://idp/token", "c", "s") + with pytest.raises(CatalogAuthError): + provider.get_token() + + @patch("requests.post") + def test_missing_access_token_raises(self, mock_post): + mock_post.return_value = Mock(ok=True, json=lambda: {}) + provider = OAuthClientCredentialsTokenProvider("http://idp/token", "c", "s") + with pytest.raises(CatalogAuthError, match="access_token"): + provider.get_token() + + @patch("requests.post") + def test_session_attaches_fresh_bearer_token(self, mock_post): + mock_post.return_value = Mock( + ok=True, json=lambda: {"access_token": "abc", "expires_in": 3600} + ) + client = PolarisClient(ENDPOINT, "w", client_id="cid", client_secret="sec") + prepared = client._session.prepare_request( + requests.Request("GET", f"{ENDPOINT}/v1/w/namespaces") + ) + assert prepared.headers["Authorization"] == "Bearer abc" + + +class TestPolarisClientProperties: + @patch("requests.Session.post") + def test_set_table_properties_request(self, mock_post): + mock_post.return_value = Mock(ok=True) + client = PolarisClient(ENDPOINT, "my_catalog", token="tok") + assert client.set_table_properties("ml", "t", {"feast.project": "p"}) is True + + url = mock_post.call_args.args[0] + assert url == f"{ENDPOINT}/v1/my_catalog/namespaces/ml/tables/t" + assert mock_post.call_args.kwargs["json"] == { + "requirements": [], + "updates": [ + {"action": "set-properties", "updates": {"feast.project": "p"}} + ], + } + + @patch("requests.Session.post") + def test_set_table_properties_noop_when_empty(self, mock_post): + client = PolarisClient(ENDPOINT, "w", token="tok") + assert client.set_table_properties("ml", "t", {}) is True + mock_post.assert_not_called() + + @patch("requests.Session.post") + def test_remove_table_properties_request(self, mock_post): + mock_post.return_value = Mock(ok=True) + client = PolarisClient(ENDPOINT, "w", token="tok") + assert client.remove_table_properties("ml", "t", ["feast.deleted_at"]) + assert mock_post.call_args.kwargs["json"]["updates"] == [ + {"action": "remove-properties", "removals": ["feast.deleted_at"]} + ] + + @patch("requests.Session.post") + def test_http_failure_returns_false_without_raising(self, mock_post): + mock_post.return_value = Mock(ok=False, status_code=403, text="forbidden") + client = PolarisClient(ENDPOINT, "w", token="tok") + assert client.set_table_properties("ml", "t", {"k": "v"}) is False + + @patch("requests.Session.post") + def test_connection_failure_returns_false(self, mock_post): + mock_post.side_effect = requests.exceptions.ConnectionError("down") + client = PolarisClient(ENDPOINT, "w", token="tok") + assert client.set_table_properties("ml", "t", {"k": "v"}) is False + + +def _feature_view(source, **kwargs) -> FeatureView: + return FeatureView( + name="driver_stats", + entities=[ + Entity( + name="driver_id", join_keys=["driver_id"], value_type=ValueType.INT64 + ) + ], + schema=[Field(name="conv_rate", dtype=Float32)], + source=source, + **kwargs, + ) + + +class TestPolarisProvider: + def test_build_properties(self): + view = _feature_view( + _source(), owner="o@example.com", description="stats", tags={"team": "ml"} + ) + assert PolarisProvider.build_properties("proj", view) == { + "feast.is_feature_table": "true", + "feast.project": "proj", + "feast.feature_view": "driver_stats", + "feast.owner": "o@example.com", + "feast.description": "stats", + "feast.tag.team": "ml", + } + + def test_sync_sets_properties_for_polaris_views_only(self): + polaris_source = _source() + client = MagicMock() + view = _feature_view(polaris_source) + other = FeatureView( + name="other", + entities=[ + Entity( + name="driver_id", + join_keys=["driver_id"], + value_type=ValueType.INT64, + ) + ], + schema=[Field(name="x", dtype=Float32)], + source=IcebergSource( + endpoint="http://h", + warehouse="w", + namespace="n", + table="t", + timestamp_field="ts", + ), + ) + provider = object.__new__(PolarisProvider) + with patch.object(PolarisSource, "get_catalog_client", return_value=client): + provider._sync_polaris_properties("proj", [view, other]) + + client.set_table_properties.assert_called_once() + args = client.set_table_properties.call_args.args + assert args[:2] == ("ml", "driver_stats") + assert args[2]["feast.project"] == "proj" + + def test_sync_respects_sync_metadata_flag(self): + client = MagicMock() + view = _feature_view(_source(sync_metadata=False)) + provider = object.__new__(PolarisProvider) + with patch.object(PolarisSource, "get_catalog_client", return_value=client): + provider._sync_polaris_properties("proj", [view]) + client.set_table_properties.assert_not_called() + + def test_deleted_view_marks_table_without_dropping(self): + client = MagicMock() + view = _feature_view(_source()) + provider = object.__new__(PolarisProvider) + with patch.object(PolarisSource, "get_catalog_client", return_value=client): + provider._mark_polaris_tables_removed([view]) + + props = client.set_table_properties.call_args.args[2] + assert props["feast.is_feature_table"] == "false" + assert "feast.deleted_at" in props + assert not client.method_calls or all( + "drop" not in c[0] for c in client.method_calls + ) + + def test_registered_in_provider_map(self): + from feast.infra.provider import PROVIDERS_CLASS_FOR_TYPE + + assert PROVIDERS_CLASS_FOR_TYPE["polaris"].endswith( + "polaris_provider.PolarisProvider" + ) 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 119e67300a5..627e8caced3 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 @@ -1267,25 +1267,7 @@ def _load_pyiceberg_table(data_source: "DataSource"): ) assert isinstance(data_source, IcebergSource) - fqn = f"{data_source.namespace}.{data_source.iceberg_table}" - - if data_source.catalog_type != "rest": - catalog_client = data_source.get_catalog_client() - return catalog_client.load_table(fqn) - - from pyiceberg.catalog import load_catalog - - catalog_config: Dict[str, str] = { - "type": "rest", - "uri": data_source.endpoint, - "warehouse": data_source.warehouse, - } - if data_source.token_env_var: - token = os.environ.get(data_source.token_env_var) - if token: - catalog_config["token"] = token - catalog = load_catalog(data_source.catalog_name, **catalog_config) - return catalog.load_table(fqn) + return data_source.load_pyiceberg_table() def _iceberg_offline_write_batch( diff --git a/sdk/python/feast/infra/offline_stores/duckdb.py b/sdk/python/feast/infra/offline_stores/duckdb.py index 8672f965fe6..541abcc351f 100644 --- a/sdk/python/feast/infra/offline_stores/duckdb.py +++ b/sdk/python/feast/infra/offline_stores/duckdb.py @@ -1,5 +1,6 @@ from __future__ import annotations +import hashlib import json import os from datetime import date, datetime, timezone @@ -140,25 +141,7 @@ def _read_iceberg_catalog_source(data_source: "IcebergSource", repo_path: str) - Loads the table through the configured catalog (REST, SQL, Hive, etc.) and scans it to Arrow. """ - fqn = f"{data_source.namespace}.{data_source.iceberg_table}" - - if data_source.catalog_type != "rest": - catalog_client = data_source.get_catalog_client() - table = catalog_client.load_table(fqn) - else: - from pyiceberg.catalog import load_catalog - - catalog_config: dict = { - "type": "rest", - "uri": data_source.endpoint, - "warehouse": data_source.warehouse, - } - if data_source.token_env_var: - token = os.environ.get(data_source.token_env_var) - if token: - catalog_config["token"] = token - catalog = load_catalog(data_source.catalog_name, **catalog_config) - table = catalog.load_table(fqn) + table = data_source.load_pyiceberg_table() return ibis.memtable(table.scan().to_arrow()) @@ -170,25 +153,7 @@ def _load_pyiceberg_table_for_write(data_source: "IcebergSource"): but returns the raw PyIceberg table instead of an ibis memtable, so the caller can use ``table.append()`` or ``table.overwrite()`` for writes. """ - fqn = f"{data_source.namespace}.{data_source.iceberg_table}" - - if data_source.catalog_type != "rest": - catalog_client = data_source.get_catalog_client() - return catalog_client.load_table(fqn) - - from pyiceberg.catalog import load_catalog - - catalog_config: dict = { - "type": "rest", - "uri": data_source.endpoint, - "warehouse": data_source.warehouse, - } - if data_source.token_env_var: - token = os.environ.get(data_source.token_env_var) - if token: - catalog_config["token"] = token - catalog = load_catalog(data_source.catalog_name, **catalog_config) - return catalog.load_table(fqn) + return data_source.load_pyiceberg_table() def _write_iceberg_data_source( @@ -229,8 +194,24 @@ def _read_lance_source(data_source: "LanceSource") -> Table: is honoured before any data crosses into ibis. The declared-schema contract is enforced separately, by ``validate_lance_source_schema``, because the reader callback is handed a ``DataSource`` without its feature view. + + The dataset is registered with DuckDB as a lazy scan rather than read into + memory. A ``LanceDataset`` is a ``pyarrow.dataset.Dataset``, so DuckDB pushes + the query's column projection into the Lance scanner: a retrieval that needs + a few scalar features never reads the embedding columns stored beside them. + + The view is named after the dataset's location and resolved version. An + opened dataset is a snapshot, so the name always denotes the same data and + re-registering it is harmless, while a job built before a later write keeps + reading the version it was built against. """ - return ibis.memtable(data_source.to_arrow()) + dataset = data_source.get_dataset() + digest = hashlib.sha256(f"{dataset.uri}@{dataset.version}".encode()).hexdigest() + view_name = f"__feast_lance_{digest[:16]}" + + backend = ibis.get_backend() + backend.con.register(view_name, dataset) + return backend.table(view_name) def _validate_lance_feature_views(feature_views: List[FeatureView]) -> None: diff --git a/sdk/python/feast/infra/provider.py b/sdk/python/feast/infra/provider.py index 2af6f84f10f..23dfd7ea93b 100644 --- a/sdk/python/feast/infra/provider.py +++ b/sdk/python/feast/infra/provider.py @@ -46,6 +46,7 @@ "azure": "feast.infra.passthrough_provider.PassthroughProvider", "chronon": "feast.infra.chronon_provider.ChrononProvider", "unity_catalog": "feast.infra.data_sources.contrib.iceberg_catalog.uc_provider.UnityCatalogProvider", + "polaris": "feast.infra.data_sources.contrib.iceberg_catalog.polaris_provider.PolarisProvider", } diff --git a/sdk/python/feast/lineage/registry_lineage.py b/sdk/python/feast/lineage/registry_lineage.py index 0f4c2ac8c83..44202cd2d47 100644 --- a/sdk/python/feast/lineage/registry_lineage.py +++ b/sdk/python/feast/lineage/registry_lineage.py @@ -228,6 +228,27 @@ def _parse_direct_relationships(self, registry: Registry) -> List[EntityRelation ) relationships.append(rel) + # Stream source -> FeatureView relationships + if ( + hasattr(feature_view.spec, "stream_source") + and feature_view.spec.stream_source + ): + stream_source_name = getattr( + feature_view.spec.stream_source, "name", None + ) + if stream_source_name: + relationships.append( + EntityRelation( + source=EntityReference( + FeastObjectType.DATA_SOURCE, stream_source_name + ), + target=EntityReference( + FeastObjectType.FEATURE_VIEW, + feature_view.spec.name, + ), + ) + ) + # Batch source relationship if ( hasattr(feature_view.spec, "batch_source") @@ -455,6 +476,72 @@ def _parse_direct_relationships(self, registry: Registry) -> List[EntityRelation ) ) + # Upstream FeatureView -> DataSource (PushSource) relationships + push_sources = { + ds.name: ds + for ds in registry.data_sources + if hasattr(ds, "name") + and ds.name + and hasattr(ds, "push_options") + and ds.push_options + } + for fv in registry.feature_views: + if ( + hasattr(fv, "spec") + and fv.spec + and hasattr(fv.spec, "stream_source") + and fv.spec.stream_source + and hasattr(fv.spec.stream_source, "name") + and fv.spec.stream_source.name + and hasattr(fv.spec.stream_source, "push_options") + and fv.spec.stream_source.push_options + ): + push_sources[fv.spec.stream_source.name] = fv.spec.stream_source + for sfv in registry.stream_feature_views: + if ( + hasattr(sfv, "spec") + and sfv.spec + and hasattr(sfv.spec, "stream_source") + and sfv.spec.stream_source + and hasattr(sfv.spec.stream_source, "name") + and sfv.spec.stream_source.name + and hasattr(sfv.spec.stream_source, "push_options") + and sfv.spec.stream_source.push_options + ): + push_sources[sfv.spec.stream_source.name] = sfv.spec.stream_source + for lv in registry.label_views: + if ( + hasattr(lv, "spec") + and lv.spec + and hasattr(lv.spec, "source") + and lv.spec.source + and hasattr(lv.spec.source, "name") + and lv.spec.source.name + and hasattr(lv.spec.source, "push_options") + and lv.spec.source.push_options + ): + push_sources[lv.spec.source.name] = lv.spec.source + + for ds in push_sources.values(): + if ( + not hasattr(ds, "push_options") + or not hasattr(ds.push_options, "upstream_feature_view_names") + or not ds.push_options.upstream_feature_view_names + ): + continue + for upstream_fv in ds.push_options.upstream_feature_view_names: + source_type = ( + FeastObjectType.LABEL_VIEW + if upstream_fv in label_view_names + else FeastObjectType.FEATURE_VIEW + ) + relationships.append( + EntityRelation( + source=EntityReference(source_type, upstream_fv), + target=EntityReference(FeastObjectType.DATA_SOURCE, ds.name), + ) + ) + # SavedDataset relationships ds_location_index = _build_datasource_location_index(registry) diff --git a/sdk/python/feast/metrics.py b/sdk/python/feast/metrics.py index e852b3d5c58..85b927ee198 100644 --- a/sdk/python/feast/metrics.py +++ b/sdk/python/feast/metrics.py @@ -681,6 +681,27 @@ def _shutdown(signum, frame): httpd.serve_forever() +# The metrics-server process started by start_metrics_server() in this process +# (the Gunicorn master), and whether the at-fork hook below is registered. +_metrics_process: Optional[multiprocessing.Process] = None +_metrics_at_fork_registered: bool = False + + +def _forget_metrics_process_in_child() -> None: + """Drop the inherited metrics-server process from multiprocessing's registry. + + Runs in every child forked from the process that started the metrics + server (for example each Gunicorn worker). A forked child inherits + multiprocessing's registry of child processes, and at interpreter exit + ``multiprocessing.util._exit_function`` terminates every daemon process in + it. Without this, the first worker exit would SIGTERM the master's metrics + server and ``/metrics`` would stop responding. The child does not own the + process; the parent still terminates it on its own exit. + """ + if _metrics_process is not None: + multiprocessing.process._children.discard(_metrics_process) # type: ignore[attr-defined] + + def start_metrics_server( store: "FeatureStore", port: int = 8000, @@ -714,7 +735,7 @@ def start_metrics_server( workers — the ``post_worker_init`` hook starts per-worker freshness monitoring instead. """ - global _config + global _config, _metrics_process, _metrics_at_fork_registered if metrics_config is not None: _config = metrics_config @@ -739,6 +760,10 @@ def start_metrics_server( name="feast-metrics-server", ) metrics_proc.start() + _metrics_process = metrics_proc + if not _metrics_at_fork_registered and hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_forget_metrics_process_in_child) + _metrics_at_fork_registered = True if ready_event.wait(timeout=5): logger.info( diff --git a/sdk/python/feast/protos/feast/core/DataSource_pb2.py b/sdk/python/feast/protos/feast/core/DataSource_pb2.py index 3ad9ab1807b..f694c357264 100644 --- a/sdk/python/feast/protos/feast/core/DataSource_pb2.py +++ b/sdk/python/feast/protos/feast/core/DataSource_pb2.py @@ -19,7 +19,7 @@ from feast.protos.feast.core import Feature_pb2 as feast_dot_core_dot_Feature__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x66\x65\x61st/core/DataSource.proto\x12\nfeast.core\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1b\x66\x65\x61st/core/DataFormat.proto\x1a\x17\x66\x65\x61st/types/Value.proto\x1a\x18\x66\x65\x61st/core/Feature.proto\"\xba\x18\n\nDataSource\x12\x0c\n\x04name\x18\x14 \x01(\t\x12\x0f\n\x07project\x18\x15 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x17 \x01(\t\x12.\n\x04tags\x18\x18 \x03(\x0b\x32 .feast.core.DataSource.TagsEntry\x12\r\n\x05owner\x18\x19 \x01(\t\x12/\n\x04type\x18\x01 \x01(\x0e\x32!.feast.core.DataSource.SourceType\x12?\n\rfield_mapping\x18\x02 \x03(\x0b\x32(.feast.core.DataSource.FieldMappingEntry\x12\x17\n\x0ftimestamp_field\x18\x03 \x01(\t\x12\x1d\n\x15\x64\x61te_partition_column\x18\x04 \x01(\t\x12 \n\x18\x63reated_timestamp_column\x18\x05 \x01(\t\x12\x1c\n\x14timestamp_field_type\x18\x1c \x01(\t\x12\x1e\n\x16\x64\x61ta_source_class_type\x18\x11 \x01(\t\x12,\n\x0c\x62\x61tch_source\x18\x1a \x01(\x0b\x32\x16.feast.core.DataSource\x12/\n\x04meta\x18\x32 \x01(\x0b\x32!.feast.core.DataSource.SourceMeta\x12:\n\x0c\x66ile_options\x18\x0b \x01(\x0b\x32\".feast.core.DataSource.FileOptionsH\x00\x12\x42\n\x10\x62igquery_options\x18\x0c \x01(\x0b\x32&.feast.core.DataSource.BigQueryOptionsH\x00\x12<\n\rkafka_options\x18\r \x01(\x0b\x32#.feast.core.DataSource.KafkaOptionsH\x00\x12@\n\x0fkinesis_options\x18\x0e \x01(\x0b\x32%.feast.core.DataSource.KinesisOptionsH\x00\x12\x42\n\x10redshift_options\x18\x0f \x01(\x0b\x32&.feast.core.DataSource.RedshiftOptionsH\x00\x12I\n\x14request_data_options\x18\x12 \x01(\x0b\x32).feast.core.DataSource.RequestDataOptionsH\x00\x12\x44\n\x0e\x63ustom_options\x18\x10 \x01(\x0b\x32*.feast.core.DataSource.CustomSourceOptionsH\x00\x12\x44\n\x11snowflake_options\x18\x13 \x01(\x0b\x32\'.feast.core.DataSource.SnowflakeOptionsH\x00\x12:\n\x0cpush_options\x18\x16 \x01(\x0b\x32\".feast.core.DataSource.PushOptionsH\x00\x12<\n\rspark_options\x18\x1b \x01(\x0b\x32#.feast.core.DataSource.SparkOptionsH\x00\x12<\n\rtrino_options\x18\x1e \x01(\x0b\x32#.feast.core.DataSource.TrinoOptionsH\x00\x12>\n\x0e\x61thena_options\x18# \x01(\x0b\x32$.feast.core.DataSource.AthenaOptionsH\x00\x1a+\n\tTagsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x1a\x33\n\x11\x46ieldMappingEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x1a\xf5\x01\n\nSourceMeta\x12:\n\x16\x65\x61rliestEventTimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x38\n\x14latestEventTimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x35\n\x11\x63reated_timestamp\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12:\n\x16last_updated_timestamp\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x1a\x65\n\x0b\x46ileOptions\x12+\n\x0b\x66ile_format\x18\x01 \x01(\x0b\x32\x16.feast.core.FileFormat\x12\x0b\n\x03uri\x18\x02 \x01(\t\x12\x1c\n\x14s3_endpoint_override\x18\x03 \x01(\t\x1a/\n\x0f\x42igQueryOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x1a,\n\x0cTrinoOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x1a\xae\x01\n\x0cKafkaOptions\x12\x1f\n\x17kafka_bootstrap_servers\x18\x01 \x01(\t\x12\r\n\x05topic\x18\x02 \x01(\t\x12\x30\n\x0emessage_format\x18\x03 \x01(\x0b\x32\x18.feast.core.StreamFormat\x12<\n\x19watermark_delay_threshold\x18\x04 \x01(\x0b\x32\x19.google.protobuf.Duration\x1a\x66\n\x0eKinesisOptions\x12\x0e\n\x06region\x18\x01 \x01(\t\x12\x13\n\x0bstream_name\x18\x02 \x01(\t\x12/\n\rrecord_format\x18\x03 \x01(\x0b\x32\x18.feast.core.StreamFormat\x1aQ\n\x0fRedshiftOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x12\x0e\n\x06schema\x18\x03 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x04 \x01(\t\x1aT\n\rAthenaOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x03 \x01(\t\x12\x13\n\x0b\x64\x61ta_source\x18\x04 \x01(\t\x1aX\n\x10SnowflakeOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x12\x0e\n\x06schema\x18\x03 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x04 \x01(\tJ\x04\x08\x05\x10\x06\x1a\xa4\x01\n\x0cSparkOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x13\n\x0b\x66ile_format\x18\x04 \x01(\t\x12$\n\x1c\x64\x61te_partition_column_format\x18\x05 \x01(\t\x12-\n\x0ctable_format\x18\x06 \x01(\x0b\x32\x17.feast.core.TableFormat\x1a,\n\x13\x43ustomSourceOptions\x12\x15\n\rconfiguration\x18\x01 \x01(\x0c\x1a\xf7\x01\n\x12RequestDataOptions\x12Z\n\x11\x64\x65precated_schema\x18\x02 \x03(\x0b\x32?.feast.core.DataSource.RequestDataOptions.DeprecatedSchemaEntry\x12)\n\x06schema\x18\x03 \x03(\x0b\x32\x19.feast.core.FeatureSpecV2\x1aT\n\x15\x44\x65precatedSchemaEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12*\n\x05value\x18\x02 \x01(\x0e\x32\x1b.feast.types.ValueType.Enum:\x02\x38\x01J\x04\x08\x01\x10\x02\x1a\x13\n\x0bPushOptionsJ\x04\x08\x01\x10\x02\"\x8b\x02\n\nSourceType\x12\x0b\n\x07INVALID\x10\x00\x12\x0e\n\nBATCH_FILE\x10\x01\x12\x13\n\x0f\x42\x41TCH_SNOWFLAKE\x10\x08\x12\x12\n\x0e\x42\x41TCH_BIGQUERY\x10\x02\x12\x12\n\x0e\x42\x41TCH_REDSHIFT\x10\x05\x12\x10\n\x0cSTREAM_KAFKA\x10\x03\x12\x12\n\x0eSTREAM_KINESIS\x10\x04\x12\x11\n\rCUSTOM_SOURCE\x10\x06\x12\x12\n\x0eREQUEST_SOURCE\x10\x07\x12\x0f\n\x0bPUSH_SOURCE\x10\t\x12\x0f\n\x0b\x42\x41TCH_TRINO\x10\n\x12\x0f\n\x0b\x42\x41TCH_SPARK\x10\x0b\x12\x10\n\x0c\x42\x41TCH_ATHENA\x10\x0c\x12\x11\n\rBATCH_ICEBERG\x10\rB\t\n\x07optionsJ\x04\x08\x06\x10\x0b\"=\n\x0e\x44\x61taSourceList\x12+\n\x0b\x64\x61tasources\x18\x01 \x03(\x0b\x32\x16.feast.core.DataSourceBT\n\x10\x66\x65\x61st.proto.coreB\x0f\x44\x61taSourceProtoZ/github.com/feast-dev/feast/go/protos/feast/coreb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1b\x66\x65\x61st/core/DataSource.proto\x12\nfeast.core\x1a\x1egoogle/protobuf/duration.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1b\x66\x65\x61st/core/DataFormat.proto\x1a\x17\x66\x65\x61st/types/Value.proto\x1a\x18\x66\x65\x61st/core/Feature.proto\"\xd4\x01\n\rConnectionRef\x12\x10\n\x08provider\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x11\n\tnamespace\x18\x03 \x01(\t\x12\x17\n\x0f\x63onnection_type\x18\x04 \x01(\t\x12\x11\n\tauth_type\x18\x05 \x01(\t\x12\x35\n\x06params\x18\x06 \x03(\x0b\x32%.feast.core.ConnectionRef.ParamsEntry\x1a-\n\x0bParamsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x92\x19\n\nDataSource\x12\x0c\n\x04name\x18\x14 \x01(\t\x12\x0f\n\x07project\x18\x15 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x17 \x01(\t\x12.\n\x04tags\x18\x18 \x03(\x0b\x32 .feast.core.DataSource.TagsEntry\x12\r\n\x05owner\x18\x19 \x01(\t\x12/\n\x04type\x18\x01 \x01(\x0e\x32!.feast.core.DataSource.SourceType\x12?\n\rfield_mapping\x18\x02 \x03(\x0b\x32(.feast.core.DataSource.FieldMappingEntry\x12\x17\n\x0ftimestamp_field\x18\x03 \x01(\t\x12\x1d\n\x15\x64\x61te_partition_column\x18\x04 \x01(\t\x12 \n\x18\x63reated_timestamp_column\x18\x05 \x01(\t\x12\x1c\n\x14timestamp_field_type\x18\x1c \x01(\t\x12\x1e\n\x16\x64\x61ta_source_class_type\x18\x11 \x01(\t\x12,\n\x0c\x62\x61tch_source\x18\x1a \x01(\x0b\x32\x16.feast.core.DataSource\x12\x31\n\x0e\x63onnection_ref\x18\x1d \x01(\x0b\x32\x19.feast.core.ConnectionRef\x12/\n\x04meta\x18\x32 \x01(\x0b\x32!.feast.core.DataSource.SourceMeta\x12:\n\x0c\x66ile_options\x18\x0b \x01(\x0b\x32\".feast.core.DataSource.FileOptionsH\x00\x12\x42\n\x10\x62igquery_options\x18\x0c \x01(\x0b\x32&.feast.core.DataSource.BigQueryOptionsH\x00\x12<\n\rkafka_options\x18\r \x01(\x0b\x32#.feast.core.DataSource.KafkaOptionsH\x00\x12@\n\x0fkinesis_options\x18\x0e \x01(\x0b\x32%.feast.core.DataSource.KinesisOptionsH\x00\x12\x42\n\x10redshift_options\x18\x0f \x01(\x0b\x32&.feast.core.DataSource.RedshiftOptionsH\x00\x12I\n\x14request_data_options\x18\x12 \x01(\x0b\x32).feast.core.DataSource.RequestDataOptionsH\x00\x12\x44\n\x0e\x63ustom_options\x18\x10 \x01(\x0b\x32*.feast.core.DataSource.CustomSourceOptionsH\x00\x12\x44\n\x11snowflake_options\x18\x13 \x01(\x0b\x32\'.feast.core.DataSource.SnowflakeOptionsH\x00\x12:\n\x0cpush_options\x18\x16 \x01(\x0b\x32\".feast.core.DataSource.PushOptionsH\x00\x12<\n\rspark_options\x18\x1b \x01(\x0b\x32#.feast.core.DataSource.SparkOptionsH\x00\x12<\n\rtrino_options\x18\x1e \x01(\x0b\x32#.feast.core.DataSource.TrinoOptionsH\x00\x12>\n\x0e\x61thena_options\x18# \x01(\x0b\x32$.feast.core.DataSource.AthenaOptionsH\x00\x1a+\n\tTagsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x1a\x33\n\x11\x46ieldMappingEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\x1a\xf5\x01\n\nSourceMeta\x12:\n\x16\x65\x61rliestEventTimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x38\n\x14latestEventTimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x35\n\x11\x63reated_timestamp\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12:\n\x16last_updated_timestamp\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x1a\x65\n\x0b\x46ileOptions\x12+\n\x0b\x66ile_format\x18\x01 \x01(\x0b\x32\x16.feast.core.FileFormat\x12\x0b\n\x03uri\x18\x02 \x01(\t\x12\x1c\n\x14s3_endpoint_override\x18\x03 \x01(\t\x1a/\n\x0f\x42igQueryOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x1a,\n\x0cTrinoOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x1a\xae\x01\n\x0cKafkaOptions\x12\x1f\n\x17kafka_bootstrap_servers\x18\x01 \x01(\t\x12\r\n\x05topic\x18\x02 \x01(\t\x12\x30\n\x0emessage_format\x18\x03 \x01(\x0b\x32\x18.feast.core.StreamFormat\x12<\n\x19watermark_delay_threshold\x18\x04 \x01(\x0b\x32\x19.google.protobuf.Duration\x1a\x66\n\x0eKinesisOptions\x12\x0e\n\x06region\x18\x01 \x01(\t\x12\x13\n\x0bstream_name\x18\x02 \x01(\t\x12/\n\rrecord_format\x18\x03 \x01(\x0b\x32\x18.feast.core.StreamFormat\x1aQ\n\x0fRedshiftOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x12\x0e\n\x06schema\x18\x03 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x04 \x01(\t\x1aT\n\rAthenaOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x03 \x01(\t\x12\x13\n\x0b\x64\x61ta_source\x18\x04 \x01(\t\x1aX\n\x10SnowflakeOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x12\x0e\n\x06schema\x18\x03 \x01(\t\x12\x10\n\x08\x64\x61tabase\x18\x04 \x01(\tJ\x04\x08\x05\x10\x06\x1a\xa4\x01\n\x0cSparkOptions\x12\r\n\x05table\x18\x01 \x01(\t\x12\r\n\x05query\x18\x02 \x01(\t\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x13\n\x0b\x66ile_format\x18\x04 \x01(\t\x12$\n\x1c\x64\x61te_partition_column_format\x18\x05 \x01(\t\x12-\n\x0ctable_format\x18\x06 \x01(\x0b\x32\x17.feast.core.TableFormat\x1a,\n\x13\x43ustomSourceOptions\x12\x15\n\rconfiguration\x18\x01 \x01(\x0c\x1a\xf7\x01\n\x12RequestDataOptions\x12Z\n\x11\x64\x65precated_schema\x18\x02 \x03(\x0b\x32?.feast.core.DataSource.RequestDataOptions.DeprecatedSchemaEntry\x12)\n\x06schema\x18\x03 \x03(\x0b\x32\x19.feast.core.FeatureSpecV2\x1aT\n\x15\x44\x65precatedSchemaEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12*\n\x05value\x18\x02 \x01(\x0e\x32\x1b.feast.types.ValueType.Enum:\x02\x38\x01J\x04\x08\x01\x10\x02\x1a\x38\n\x0bPushOptions\x12#\n\x1bupstream_feature_view_names\x18\x02 \x03(\tJ\x04\x08\x01\x10\x02\"\x8b\x02\n\nSourceType\x12\x0b\n\x07INVALID\x10\x00\x12\x0e\n\nBATCH_FILE\x10\x01\x12\x13\n\x0f\x42\x41TCH_SNOWFLAKE\x10\x08\x12\x12\n\x0e\x42\x41TCH_BIGQUERY\x10\x02\x12\x12\n\x0e\x42\x41TCH_REDSHIFT\x10\x05\x12\x10\n\x0cSTREAM_KAFKA\x10\x03\x12\x12\n\x0eSTREAM_KINESIS\x10\x04\x12\x11\n\rCUSTOM_SOURCE\x10\x06\x12\x12\n\x0eREQUEST_SOURCE\x10\x07\x12\x0f\n\x0bPUSH_SOURCE\x10\t\x12\x0f\n\x0b\x42\x41TCH_TRINO\x10\n\x12\x0f\n\x0b\x42\x41TCH_SPARK\x10\x0b\x12\x10\n\x0c\x42\x41TCH_ATHENA\x10\x0c\x12\x11\n\rBATCH_ICEBERG\x10\rB\t\n\x07optionsJ\x04\x08\x06\x10\x0b\"=\n\x0e\x44\x61taSourceList\x12+\n\x0b\x64\x61tasources\x18\x01 \x03(\x0b\x32\x16.feast.core.DataSourceBT\n\x10\x66\x65\x61st.proto.coreB\x0f\x44\x61taSourceProtoZ/github.com/feast-dev/feast/go/protos/feast/coreb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -27,48 +27,54 @@ if _descriptor._USE_C_DESCRIPTORS == False: _globals['DESCRIPTOR']._options = None _globals['DESCRIPTOR']._serialized_options = b'\n\020feast.proto.coreB\017DataSourceProtoZ/github.com/feast-dev/feast/go/protos/feast/core' + _globals['_CONNECTIONREF_PARAMSENTRY']._options = None + _globals['_CONNECTIONREF_PARAMSENTRY']._serialized_options = b'8\001' _globals['_DATASOURCE_TAGSENTRY']._options = None _globals['_DATASOURCE_TAGSENTRY']._serialized_options = b'8\001' _globals['_DATASOURCE_FIELDMAPPINGENTRY']._options = None _globals['_DATASOURCE_FIELDMAPPINGENTRY']._serialized_options = b'8\001' _globals['_DATASOURCE_REQUESTDATAOPTIONS_DEPRECATEDSCHEMAENTRY']._options = None _globals['_DATASOURCE_REQUESTDATAOPTIONS_DEPRECATEDSCHEMAENTRY']._serialized_options = b'8\001' - _globals['_DATASOURCE']._serialized_start=189 - _globals['_DATASOURCE']._serialized_end=3319 - _globals['_DATASOURCE_TAGSENTRY']._serialized_start=1466 - _globals['_DATASOURCE_TAGSENTRY']._serialized_end=1509 - _globals['_DATASOURCE_FIELDMAPPINGENTRY']._serialized_start=1511 - _globals['_DATASOURCE_FIELDMAPPINGENTRY']._serialized_end=1562 - _globals['_DATASOURCE_SOURCEMETA']._serialized_start=1565 - _globals['_DATASOURCE_SOURCEMETA']._serialized_end=1810 - _globals['_DATASOURCE_FILEOPTIONS']._serialized_start=1812 - _globals['_DATASOURCE_FILEOPTIONS']._serialized_end=1913 - _globals['_DATASOURCE_BIGQUERYOPTIONS']._serialized_start=1915 - _globals['_DATASOURCE_BIGQUERYOPTIONS']._serialized_end=1962 - _globals['_DATASOURCE_TRINOOPTIONS']._serialized_start=1964 - _globals['_DATASOURCE_TRINOOPTIONS']._serialized_end=2008 - _globals['_DATASOURCE_KAFKAOPTIONS']._serialized_start=2011 - _globals['_DATASOURCE_KAFKAOPTIONS']._serialized_end=2185 - _globals['_DATASOURCE_KINESISOPTIONS']._serialized_start=2187 - _globals['_DATASOURCE_KINESISOPTIONS']._serialized_end=2289 - _globals['_DATASOURCE_REDSHIFTOPTIONS']._serialized_start=2291 - _globals['_DATASOURCE_REDSHIFTOPTIONS']._serialized_end=2372 - _globals['_DATASOURCE_ATHENAOPTIONS']._serialized_start=2374 - _globals['_DATASOURCE_ATHENAOPTIONS']._serialized_end=2458 - _globals['_DATASOURCE_SNOWFLAKEOPTIONS']._serialized_start=2460 - _globals['_DATASOURCE_SNOWFLAKEOPTIONS']._serialized_end=2548 - _globals['_DATASOURCE_SPARKOPTIONS']._serialized_start=2551 - _globals['_DATASOURCE_SPARKOPTIONS']._serialized_end=2715 - _globals['_DATASOURCE_CUSTOMSOURCEOPTIONS']._serialized_start=2717 - _globals['_DATASOURCE_CUSTOMSOURCEOPTIONS']._serialized_end=2761 - _globals['_DATASOURCE_REQUESTDATAOPTIONS']._serialized_start=2764 - _globals['_DATASOURCE_REQUESTDATAOPTIONS']._serialized_end=3011 - _globals['_DATASOURCE_REQUESTDATAOPTIONS_DEPRECATEDSCHEMAENTRY']._serialized_start=2921 - _globals['_DATASOURCE_REQUESTDATAOPTIONS_DEPRECATEDSCHEMAENTRY']._serialized_end=3005 - _globals['_DATASOURCE_PUSHOPTIONS']._serialized_start=3013 - _globals['_DATASOURCE_PUSHOPTIONS']._serialized_end=3032 - _globals['_DATASOURCE_SOURCETYPE']._serialized_start=3035 - _globals['_DATASOURCE_SOURCETYPE']._serialized_end=3302 - _globals['_DATASOURCELIST']._serialized_start=3321 - _globals['_DATASOURCELIST']._serialized_end=3382 + _globals['_CONNECTIONREF']._serialized_start=189 + _globals['_CONNECTIONREF']._serialized_end=401 + _globals['_CONNECTIONREF_PARAMSENTRY']._serialized_start=356 + _globals['_CONNECTIONREF_PARAMSENTRY']._serialized_end=401 + _globals['_DATASOURCE']._serialized_start=404 + _globals['_DATASOURCE']._serialized_end=3622 + _globals['_DATASOURCE_TAGSENTRY']._serialized_start=1732 + _globals['_DATASOURCE_TAGSENTRY']._serialized_end=1775 + _globals['_DATASOURCE_FIELDMAPPINGENTRY']._serialized_start=1777 + _globals['_DATASOURCE_FIELDMAPPINGENTRY']._serialized_end=1828 + _globals['_DATASOURCE_SOURCEMETA']._serialized_start=1831 + _globals['_DATASOURCE_SOURCEMETA']._serialized_end=2076 + _globals['_DATASOURCE_FILEOPTIONS']._serialized_start=2078 + _globals['_DATASOURCE_FILEOPTIONS']._serialized_end=2179 + _globals['_DATASOURCE_BIGQUERYOPTIONS']._serialized_start=2181 + _globals['_DATASOURCE_BIGQUERYOPTIONS']._serialized_end=2228 + _globals['_DATASOURCE_TRINOOPTIONS']._serialized_start=2230 + _globals['_DATASOURCE_TRINOOPTIONS']._serialized_end=2274 + _globals['_DATASOURCE_KAFKAOPTIONS']._serialized_start=2277 + _globals['_DATASOURCE_KAFKAOPTIONS']._serialized_end=2451 + _globals['_DATASOURCE_KINESISOPTIONS']._serialized_start=2453 + _globals['_DATASOURCE_KINESISOPTIONS']._serialized_end=2555 + _globals['_DATASOURCE_REDSHIFTOPTIONS']._serialized_start=2557 + _globals['_DATASOURCE_REDSHIFTOPTIONS']._serialized_end=2638 + _globals['_DATASOURCE_ATHENAOPTIONS']._serialized_start=2640 + _globals['_DATASOURCE_ATHENAOPTIONS']._serialized_end=2724 + _globals['_DATASOURCE_SNOWFLAKEOPTIONS']._serialized_start=2726 + _globals['_DATASOURCE_SNOWFLAKEOPTIONS']._serialized_end=2814 + _globals['_DATASOURCE_SPARKOPTIONS']._serialized_start=2817 + _globals['_DATASOURCE_SPARKOPTIONS']._serialized_end=2981 + _globals['_DATASOURCE_CUSTOMSOURCEOPTIONS']._serialized_start=2983 + _globals['_DATASOURCE_CUSTOMSOURCEOPTIONS']._serialized_end=3027 + _globals['_DATASOURCE_REQUESTDATAOPTIONS']._serialized_start=3030 + _globals['_DATASOURCE_REQUESTDATAOPTIONS']._serialized_end=3277 + _globals['_DATASOURCE_REQUESTDATAOPTIONS_DEPRECATEDSCHEMAENTRY']._serialized_start=3187 + _globals['_DATASOURCE_REQUESTDATAOPTIONS_DEPRECATEDSCHEMAENTRY']._serialized_end=3271 + _globals['_DATASOURCE_PUSHOPTIONS']._serialized_start=3279 + _globals['_DATASOURCE_PUSHOPTIONS']._serialized_end=3335 + _globals['_DATASOURCE_SOURCETYPE']._serialized_start=3338 + _globals['_DATASOURCE_SOURCETYPE']._serialized_end=3605 + _globals['_DATASOURCELIST']._serialized_start=3624 + _globals['_DATASOURCELIST']._serialized_end=3685 # @@protoc_insertion_point(module_scope) diff --git a/sdk/python/feast/protos/feast/core/DataSource_pb2.pyi b/sdk/python/feast/protos/feast/core/DataSource_pb2.pyi index 0b77320af32..1ebbe1981f6 100644 --- a/sdk/python/feast/protos/feast/core/DataSource_pb2.pyi +++ b/sdk/python/feast/protos/feast/core/DataSource_pb2.pyi @@ -37,9 +37,76 @@ else: DESCRIPTOR: google.protobuf.descriptor.FileDescriptor +class ConnectionRef(google.protobuf.message.Message): + """Connection reference for a DataSource. + Combines connection type, credential resolution, and non-sensitive + connection parameters into a single reusable reference. + Allows DataSources to declare their full connection identity — which + backend to use, how to authenticate, and where to connect — independent + of the global offline_store config in feature_store.yaml. + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + class ParamsEntry(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + KEY_FIELD_NUMBER: builtins.int + VALUE_FIELD_NUMBER: builtins.int + key: builtins.str + value: builtins.str + def __init__( + self, + *, + key: builtins.str = ..., + value: builtins.str = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["key", b"key", "value", b"value"]) -> None: ... + + PROVIDER_FIELD_NUMBER: builtins.int + NAME_FIELD_NUMBER: builtins.int + NAMESPACE_FIELD_NUMBER: builtins.int + CONNECTION_TYPE_FIELD_NUMBER: builtins.int + AUTH_TYPE_FIELD_NUMBER: builtins.int + PARAMS_FIELD_NUMBER: builtins.int + provider: builtins.str + """Credential provider type: "kubernetes", "vault", "aws-secrets-manager", + "gcp-secret-manager", "azure-key-vault", "env". + """ + name: builtins.str + """Provider-specific name: K8s Secret name, Vault path, env var prefix, etc.""" + namespace: builtins.str + """Optional scope qualifier: K8s namespace, Vault mount, AWS region, etc.""" + connection_type: builtins.str + """Optional offline store class type (e.g., "snowflake.offline", "bigquery", + "spark", "iceberg-rest"). When empty, inferred from the DataSource class. + """ + auth_type: builtins.str + """Optional authentication mechanism: "secret", "oauth2", "basic", "sigv4". + Defaults to "secret" (raw key-value credentials from provider). + """ + @property + def params(self) -> google.protobuf.internal.containers.ScalarMap[builtins.str, builtins.str]: + """Optional non-sensitive connection parameters (e.g., account, database, + warehouse, endpoint URI). Keys and values are backend-specific. + """ + def __init__( + self, + *, + provider: builtins.str = ..., + name: builtins.str = ..., + namespace: builtins.str = ..., + connection_type: builtins.str = ..., + auth_type: builtins.str = ..., + params: collections.abc.Mapping[builtins.str, builtins.str] | None = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["auth_type", b"auth_type", "connection_type", b"connection_type", "name", b"name", "namespace", b"namespace", "params", b"params", "provider", b"provider"]) -> None: ... + +global___ConnectionRef = ConnectionRef + class DataSource(google.protobuf.message.Message): """Defines a Data Source that can be used source Feature data - Next available id: 29 + Next available id: 30 """ DESCRIPTOR: google.protobuf.descriptor.Descriptor @@ -457,9 +524,16 @@ class DataSource(google.protobuf.message.Message): DESCRIPTOR: google.protobuf.descriptor.Descriptor + UPSTREAM_FEATURE_VIEW_NAMES_FIELD_NUMBER: builtins.int + @property + def upstream_feature_view_names(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Names of upstream FeatureViews consumed by the process feeding this PushSource""" def __init__( self, + *, + upstream_feature_view_names: collections.abc.Iterable[builtins.str] | None = ..., ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["upstream_feature_view_names", b"upstream_feature_view_names"]) -> None: ... NAME_FIELD_NUMBER: builtins.int PROJECT_FIELD_NUMBER: builtins.int @@ -474,6 +548,7 @@ class DataSource(google.protobuf.message.Message): TIMESTAMP_FIELD_TYPE_FIELD_NUMBER: builtins.int DATA_SOURCE_CLASS_TYPE_FIELD_NUMBER: builtins.int BATCH_SOURCE_FIELD_NUMBER: builtins.int + CONNECTION_REF_FIELD_NUMBER: builtins.int META_FIELD_NUMBER: builtins.int FILE_OPTIONS_FIELD_NUMBER: builtins.int BIGQUERY_OPTIONS_FIELD_NUMBER: builtins.int @@ -523,6 +598,14 @@ class DataSource(google.protobuf.message.Message): def batch_source(self) -> global___DataSource: """Optional batch source for streaming sources for historical features and materialization.""" @property + def connection_ref(self) -> global___ConnectionRef: + """Optional connection reference for this data source. + When set, OfflineStores resolve connection type and credentials at + runtime via the registered CredentialProvider instead of using ambient + env vars or the global offline_store config in feature_store.yaml. + All fields except provider and name are optional. + """ + @property def meta(self) -> global___DataSource.SourceMeta: ... @property def file_options(self) -> global___DataSource.FileOptions: ... @@ -564,6 +647,7 @@ class DataSource(google.protobuf.message.Message): timestamp_field_type: builtins.str = ..., data_source_class_type: builtins.str = ..., batch_source: global___DataSource | None = ..., + connection_ref: global___ConnectionRef | None = ..., meta: global___DataSource.SourceMeta | None = ..., file_options: global___DataSource.FileOptions | None = ..., bigquery_options: global___DataSource.BigQueryOptions | None = ..., @@ -578,8 +662,8 @@ class DataSource(google.protobuf.message.Message): trino_options: global___DataSource.TrinoOptions | None = ..., athena_options: global___DataSource.AthenaOptions | None = ..., ) -> None: ... - def HasField(self, field_name: typing_extensions.Literal["athena_options", b"athena_options", "batch_source", b"batch_source", "bigquery_options", b"bigquery_options", "custom_options", b"custom_options", "file_options", b"file_options", "kafka_options", b"kafka_options", "kinesis_options", b"kinesis_options", "meta", b"meta", "options", b"options", "push_options", b"push_options", "redshift_options", b"redshift_options", "request_data_options", b"request_data_options", "snowflake_options", b"snowflake_options", "spark_options", b"spark_options", "trino_options", b"trino_options"]) -> builtins.bool: ... - def ClearField(self, field_name: typing_extensions.Literal["athena_options", b"athena_options", "batch_source", b"batch_source", "bigquery_options", b"bigquery_options", "created_timestamp_column", b"created_timestamp_column", "custom_options", b"custom_options", "data_source_class_type", b"data_source_class_type", "date_partition_column", b"date_partition_column", "description", b"description", "field_mapping", b"field_mapping", "file_options", b"file_options", "kafka_options", b"kafka_options", "kinesis_options", b"kinesis_options", "meta", b"meta", "name", b"name", "options", b"options", "owner", b"owner", "project", b"project", "push_options", b"push_options", "redshift_options", b"redshift_options", "request_data_options", b"request_data_options", "snowflake_options", b"snowflake_options", "spark_options", b"spark_options", "tags", b"tags", "timestamp_field", b"timestamp_field", "timestamp_field_type", b"timestamp_field_type", "trino_options", b"trino_options", "type", b"type"]) -> None: ... + def HasField(self, field_name: typing_extensions.Literal["athena_options", b"athena_options", "batch_source", b"batch_source", "bigquery_options", b"bigquery_options", "connection_ref", b"connection_ref", "custom_options", b"custom_options", "file_options", b"file_options", "kafka_options", b"kafka_options", "kinesis_options", b"kinesis_options", "meta", b"meta", "options", b"options", "push_options", b"push_options", "redshift_options", b"redshift_options", "request_data_options", b"request_data_options", "snowflake_options", b"snowflake_options", "spark_options", b"spark_options", "trino_options", b"trino_options"]) -> builtins.bool: ... + def ClearField(self, field_name: typing_extensions.Literal["athena_options", b"athena_options", "batch_source", b"batch_source", "bigquery_options", b"bigquery_options", "connection_ref", b"connection_ref", "created_timestamp_column", b"created_timestamp_column", "custom_options", b"custom_options", "data_source_class_type", b"data_source_class_type", "date_partition_column", b"date_partition_column", "description", b"description", "field_mapping", b"field_mapping", "file_options", b"file_options", "kafka_options", b"kafka_options", "kinesis_options", b"kinesis_options", "meta", b"meta", "name", b"name", "options", b"options", "owner", b"owner", "project", b"project", "push_options", b"push_options", "redshift_options", b"redshift_options", "request_data_options", b"request_data_options", "snowflake_options", b"snowflake_options", "spark_options", b"spark_options", "tags", b"tags", "timestamp_field", b"timestamp_field", "timestamp_field_type", b"timestamp_field_type", "trino_options", b"trino_options", "type", b"type"]) -> None: ... def WhichOneof(self, oneof_group: typing_extensions.Literal["options", b"options"]) -> typing_extensions.Literal["file_options", "bigquery_options", "kafka_options", "kinesis_options", "redshift_options", "request_data_options", "custom_options", "snowflake_options", "push_options", "spark_options", "trino_options", "athena_options"] | None: ... global___DataSource = DataSource diff --git a/sdk/python/feast/templates/aws/README.md b/sdk/python/feast/templates/aws/README.md index 10ae4887551..ab113b09fcd 100644 --- a/sdk/python/feast/templates/aws/README.md +++ b/sdk/python/feast/templates/aws/README.md @@ -10,7 +10,7 @@ You can run the overall workflow with `python test_workflow.py`. ## To move from this into a more production ready workflow: 1. `feature_store.yaml` points to a local file as a registry. You'll want to setup a remote file (e.g. in S3/GCS) or a -SQL registry. See [registry docs](https://docs.feast.dev/getting-started/concepts/registry) for more details. +SQL registry. See [registry docs](https://docs.feast.dev/getting-started/components/registry) for more details. 2. Setup CI/CD + dev vs staging vs prod environments to automatically update the registry as you change Feast feature definitions. See [docs](https://docs.feast.dev/how-to-guides/running-feast-in-production#1.-automatically-deploying-changes-to-your-feature-definitions). 3. (optional) Regularly scheduled materialization to power low latency feature retrieval (e.g. via Airflow). See [Batch data ingestion](https://docs.feast.dev/getting-started/concepts/data-ingestion#batch-data-ingestion) for more details. diff --git a/sdk/python/feast/templates/gcp/README.md b/sdk/python/feast/templates/gcp/README.md index 7f28fdbfa2f..3a0a9cd480e 100644 --- a/sdk/python/feast/templates/gcp/README.md +++ b/sdk/python/feast/templates/gcp/README.md @@ -10,7 +10,7 @@ You can run the overall workflow with `python test_workflow.py`. ## To move from this into a more production ready workflow: 1. `feature_store.yaml` points to a local file as a registry. You'll want to setup a remote file (e.g. in S3/GCS) or a -SQL registry. See [registry docs](https://docs.feast.dev/getting-started/concepts/registry) for more details. +SQL registry. See [registry docs](https://docs.feast.dev/getting-started/components/registry) for more details. 2. This example uses an already setup BigQuery Feast data warehouse as the [offline store](https://docs.feast.dev/getting-started/components/offline-store) to generate training data. You'll need to connect your own BigQuery instance to make this work. 3. Setup CI/CD + dev vs staging vs prod environments to automatically update the registry as you change Feast feature definitions. See [docs](https://docs.feast.dev/how-to-guides/running-feast-in-production#1.-automatically-deploying-changes-to-your-feature-definitions). diff --git a/sdk/python/tests/component/ray/test_compute.py b/sdk/python/tests/component/ray/test_compute.py index 84735e87bee..a8fe8961884 100644 --- a/sdk/python/tests/component/ray/test_compute.py +++ b/sdk/python/tests/component/ray/test_compute.py @@ -41,9 +41,10 @@ def test_ray_compute_engine_get_historical_features( fs = ray_environment.feature_store registry = fs.registry + # The udf runs on raw source rows, before aggregation. def transform_feature(df: pd.DataFrame) -> pd.DataFrame: - df["sum_conv_rate"] = df["sum_conv_rate"] * 2 - df["avg_acc_rate"] = df["avg_acc_rate"] * 2 + df["conv_rate"] = df["conv_rate"] * 2 + df["acc_rate"] = df["acc_rate"] * 2 return df driver_stats_fv = BatchFeatureView( @@ -102,11 +103,6 @@ def test_ray_compute_engine_materialize(ray_environment, feature_dataset): fs = ray_environment.feature_store registry = fs.registry - def transform_feature(df: pd.DataFrame) -> pd.DataFrame: - df["sum_conv_rate"] = df["sum_conv_rate"] * 2 - df["avg_acc_rate"] = df["avg_acc_rate"] * 2 - return df - driver_stats_fv = BatchFeatureView( name="driver_hourly_stats", entities=[driver], @@ -115,13 +111,10 @@ def transform_feature(df: pd.DataFrame) -> pd.DataFrame: Aggregation(column="conv_rate", function="sum"), Aggregation(column="acc_rate", function="avg"), ], - udf=transform_feature, - udf_string="transform_feature", ttl=timedelta(days=3), schema=[ - Field(name="conv_rate", dtype=Float32), - Field(name="acc_rate", dtype=Float32), - Field(name="avg_daily_trips", dtype=Int64), + Field(name="sum_conv_rate", dtype=Float32), + Field(name="avg_acc_rate", dtype=Float32), Field(name="driver_id", dtype=Int32), ], online=True, @@ -153,6 +146,17 @@ def tqdm_builder(length): assert len(ray_materialize_jobs) == 1 assert ray_materialize_jobs[0].status() == MaterializationJobStatus.SUCCEEDED + # driver 1001 has one row in the window: conv_rate 0.8, acc_rate 0.5 + online = fs.get_online_features( + features=[ + "driver_hourly_stats:sum_conv_rate", + "driver_hourly_stats:avg_acc_rate", + ], + entity_rows=[{"driver_id": 1001}], + ).to_dict() + assert online["sum_conv_rate"][0] == pytest.approx(0.8) + assert online["avg_acc_rate"][0] == pytest.approx(0.5) + @pytest.mark.integration @pytest.mark.xdist_group(name="ray") diff --git a/sdk/python/tests/component/ray/test_nodes.py b/sdk/python/tests/component/ray/test_nodes.py index 983b40a5ccc..3feb051a5f0 100644 --- a/sdk/python/tests/component/ray/test_nodes.py +++ b/sdk/python/tests/component/ray/test_nodes.py @@ -6,6 +6,7 @@ import pytest import ray +from feast import Entity, FeatureView, Field, FileSource, RepoConfig from feast.aggregation import Aggregation from feast.infra.compute_engines.dag.context import ColumnInfo from feast.infra.compute_engines.dag.model import DAGFormat @@ -19,13 +20,16 @@ RayJoinNode, RayReadNode, RayTransformationNode, + RayWriteNode, ) +from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig from feast.infra.ray_initializer import ( RayConfigManager, RayExecutionMode, ensure_ray_initialized, get_ray_wrapper, ) +from feast.types import Float32 class DummyInputNode(DAGNode): @@ -195,6 +199,34 @@ def test_ray_aggregation_node( assert "avg_acc_rate" in result_df.columns +def test_ray_tiled_aggregation_node_keeps_created_timestamp( + ray_session, ray_config, mock_context, sample_data +): + """Tiled aggregation must keep the latest created timestamp per entity.""" + input_value = DAGValue(data=ray.data.from_pandas(sample_data), format=DAGFormat.RAY) + node = RayAggregationNode( + name="aggregation", + aggregations=[ + Aggregation( + column="conv_rate", function="sum", time_window=timedelta(days=1) + ) + ], + group_by_keys=["driver_id"], + timestamp_col="event_timestamp", + config=ray_config, + enable_tiling=True, + hop_size=timedelta(hours=1), + created_timestamp_col="created", + ) + node.add_input(DummyInputNode("input_node", input_value)) + mock_context.node_outputs = {"input_node": input_value} + + result_df = node.execute(mock_context).data.to_pandas().set_index("driver_id") + + expected = sample_data.groupby("driver_id")["created"].max() + assert (result_df["created"] == expected.loc[result_df.index]).all() + + def test_ray_join_node(ray_session, ray_config, mock_context, sample_data, column_info): """Test RayJoinNode functionality.""" entity_data = pd.DataFrame( @@ -439,6 +471,38 @@ def test_ray_dedup_node_materialization_no_join_keys( ) +def test_ray_write_node_raises_when_online_write_fails( + ray_session, mock_context, sample_data, tmp_path +): + """A failed online store write must fail the write node, not be swallowed.""" + driver = Entity(name="driver", join_keys=["driver_id"]) + feature_view = FeatureView( + name="driver_stats", + entities=[driver], + schema=[Field(name="conv_rate", dtype=Float32)], + source=FileSource( + path=str(tmp_path / "source.parquet"), timestamp_field="event_timestamp" + ), + online=True, + offline=False, + ) + # The online table was never created, so every write raises "no such table". + mock_context.repo_config = RepoConfig( + project="test_project", + registry=str(tmp_path / "registry.db"), + provider="local", + entity_key_serialization_version=3, + online_store=SqliteOnlineStoreConfig(path=str(tmp_path / "online.db")), + ) + input_value = DAGValue(data=ray.data.from_pandas(sample_data), format=DAGFormat.RAY) + node = RayWriteNode(name="write", feature_view=feature_view) + node.add_input(DummyInputNode("input_node", input_value)) + mock_context.node_outputs = {"input_node": input_value} + + with pytest.raises(Exception, match="no such table"): + node.execute(mock_context) + + def test_ray_config_validation(): """Test Ray configuration validation.""" # Test valid configuration diff --git a/sdk/python/tests/unit/api/test_api_rest_registry.py b/sdk/python/tests/unit/api/test_api_rest_registry.py index 91b97c4406d..00651db0f3c 100644 --- a/sdk/python/tests/unit/api/test_api_rest_registry.py +++ b/sdk/python/tests/unit/api/test_api_rest_registry.py @@ -215,6 +215,26 @@ def test_feature_views_type_field_via_rest(fastapi_test_app): assert data["spec"]["name"] == "user_profile" +def test_feature_views_spec_project_via_rest(fastapi_test_app): + """Test that spec.project is populated on every feature view endpoint.""" + response = fastapi_test_app.get("/feature_views?project=demo_project") + assert response.status_code == 200 + for fv in response.json()["featureViews"]: + assert fv["spec"]["project"] == "demo_project" + + response = fastapi_test_app.get("/feature_views/user_profile?project=demo_project") + assert response.status_code == 200 + assert response.json()["spec"]["project"] == "demo_project" + + response = fastapi_test_app.get("/feature_views/all") + assert response.status_code == 200 + feature_views = response.json()["featureViews"] + assert feature_views + for fv in feature_views: + assert fv["project"] == "demo_project" + assert fv["spec"]["project"] == "demo_project" + + def test_feature_views_entity_filtering_via_rest(fastapi_test_app): """Test that feature views can be filtered by entity.""" response = fastapi_test_app.get("/feature_views?project=demo_project") diff --git a/sdk/python/tests/unit/infra/data_sources/test_lance_source.py b/sdk/python/tests/unit/infra/data_sources/test_lance_source.py index 4c68ec8db54..5c843688098 100644 --- a/sdk/python/tests/unit/infra/data_sources/test_lance_source.py +++ b/sdk/python/tests/unit/infra/data_sources/test_lance_source.py @@ -592,6 +592,94 @@ def test_pull_latest_from_a_lance_source(tmp_path): assert job.to_arrow().num_rows == 3 +def _driver_stats_with_embeddings(n_rows=3, width=8): + stats = _driver_stats(n_rows=n_rows) + embeddings = pa.FixedSizeListArray.from_arrays( + pa.array([0.5] * (n_rows * width), pa.float32()), width + ) + return stats.append_column("embedding", embeddings) + + +@pytest.fixture +def scanned_columns(monkeypatch): + """Record the column projection of every scan Lance is asked to run.""" + scans = [] + original = lance.LanceDataset.scanner + + def recording_scanner(self, *args, **kwargs): + scans.append(kwargs.get("columns")) + return original(self, *args, **kwargs) + + monkeypatch.setattr(lance.LanceDataset, "scanner", recording_scanner) + return scans + + +def test_the_duckdb_reader_is_lazy(tmp_path, scanned_columns): + uri = str(tmp_path / "stats.lance") + _write(uri, _driver_stats()) + table = _read_data_source( + LanceSource(uri=uri, timestamp_field="event_timestamp"), str(tmp_path) + ) + assert scanned_columns == [] + assert table.to_pyarrow().num_rows == 3 + assert scanned_columns + + +def test_pull_latest_does_not_scan_columns_it_does_not_select( + tmp_path, scanned_columns +): + uri = str(tmp_path / "stats.lance") + _write(uri, _driver_stats_with_embeddings()) + job = DuckDBOfflineStore.pull_latest_from_table_or_query( + config=_repo_config(tmp_path), + data_source=LanceSource(uri=uri, timestamp_field="event_timestamp"), + join_key_columns=["driver_id"], + feature_name_columns=["conv_rate"], + timestamp_field="event_timestamp", + created_timestamp_column=None, + start_date=datetime(2026, 1, 1), + end_date=datetime(2026, 1, 2), + ) + result = job.to_arrow() + assert result.num_rows == 3 + assert "embedding" not in result.column_names + assert scanned_columns + assert all( + columns is not None and "embedding" not in columns + for columns in scanned_columns + ) + + +def test_historical_retrieval_does_not_scan_columns_it_does_not_select( + tmp_path, scanned_columns +): + uri = str(tmp_path / "stats.lance") + _write(uri, _driver_stats_with_embeddings()) + source = LanceSource(uri=uri, timestamp_field="event_timestamp") + df = _historical(tmp_path, source, entity_hour=13) + # driver 1000 is the first row of the fixture. + assert df["conv_rate"].iloc[0] == pytest.approx(0.1) + assert scanned_columns + assert all( + columns is not None and "embedding" not in columns + for columns in scanned_columns + ) + + +def test_a_job_reads_the_version_it_was_built_against(tmp_path): + """A lazy scan must not pick up a write that lands before it executes.""" + uri = str(tmp_path / "stats.lance") + _write(uri, _driver_stats(n_rows=3)) + source = LanceSource(uri=uri, timestamp_field="event_timestamp") + before = _read_data_source(source, str(tmp_path)) + + _write(uri, _driver_stats(n_rows=2, base_hour=12), mode="append") + after = _read_data_source(source, str(tmp_path)) + + assert before.to_pyarrow().num_rows == 3 + assert after.to_pyarrow().num_rows == 5 + + def _pit_fixture(tmp_path, namespace=None): """One driver with two feature rows, an hour apart. diff --git a/sdk/python/tests/unit/infra/test_registry_lineage.py b/sdk/python/tests/unit/infra/test_registry_lineage.py index 5446eb2e0b6..dbdf3af54a0 100644 --- a/sdk/python/tests/unit/infra/test_registry_lineage.py +++ b/sdk/python/tests/unit/infra/test_registry_lineage.py @@ -644,3 +644,329 @@ def test_complex_entity_data_source_lineage(self): entity_name = f"entity_{i}" expected_ds = {f"data_source_{i}", "data_source_0"} assert entity_ds_connections[entity_name] == expected_ds + + def test_push_source_upstream_lineage(self): + """Test Upstream FeatureViews -> PushSource -> Target FeatureView lineage.""" + registry = Registry() + + # 1. Upstream FeatureViews + fv1_spec = FeatureViewSpec(name="user_transaction_stats") + fv1 = FeatureView(spec=fv1_spec) + + fv2_spec = FeatureViewSpec(name="user_credit_profile") + fv2 = FeatureView(spec=fv2_spec) + + # 2. PushSource + push_ds = DataSource() + push_ds.name = "risk_calc_pipeline" + push_ds.type = DataSource.SourceType.PUSH_SOURCE + push_ds.push_options.upstream_feature_view_names.extend( + ["user_transaction_stats", "user_credit_profile"] + ) + registry.data_sources.append(push_ds) + + # 3. Downstream Target FeatureView + target_fv_spec = FeatureViewSpec(name="user_risk_target_fv") + target_fv_spec.stream_source.CopyFrom(push_ds) + target_fv = FeatureView(spec=target_fv_spec) + + registry.feature_views.extend([fv1, fv2, target_fv]) + + lineage_generator = RegistryLineageGenerator() + direct_relationships, indirect_relationships = ( + lineage_generator.generate_lineage(registry) + ) + + # Verify direct relationships + expected_direct = { + ( + "featureView", + "user_transaction_stats", + "dataSource", + "risk_calc_pipeline", + ), + ("featureView", "user_credit_profile", "dataSource", "risk_calc_pipeline"), + ("dataSource", "risk_calc_pipeline", "featureView", "user_risk_target_fv"), + } + + actual_direct = { + ( + rel.source.type.value, + rel.source.name, + rel.target.type.value, + rel.target.name, + ) + for rel in direct_relationships + } + + for exp in expected_direct: + assert exp in actual_direct, ( + f"Expected relationship {exp} not found in {actual_direct}" + ) + + # Test object relationships for push source + ps_relationships = lineage_generator.get_object_relationships( + registry, "dataSource", "risk_calc_pipeline", include_indirect=False + ) + assert len(ps_relationships) == 3 + + # Test object lineage graph for upstream feature view traversing downstream + graph = lineage_generator.get_object_lineage_graph( + registry, "featureView", "user_transaction_stats", depth=2 + ) + assert "nodes" in graph + assert "featureView:user_transaction_stats" in graph["nodes"] + assert "dataSource:risk_calc_pipeline" in graph["nodes"] + assert "featureView:user_risk_target_fv" in graph["nodes"] + + def test_feature_view_stream_sources_without_upstream_views(self): + """Test FeatureView with stream sources (PushSource, KafkaSource, KinesisSource) + without upstream_feature_view_names specified draws edges from stream_source to FeatureView + and batch_source to FeatureView. + """ + registry = Registry() + + # 1. FeatureView with PushSource (no upstream_feature_view_names) + push_ds = DataSource() + push_ds.name = "user_push_source" + push_ds.type = DataSource.SourceType.PUSH_SOURCE + + push_batch_ds = DataSource() + push_batch_ds.name = "user_push_batch_source" + push_batch_ds.type = DataSource.SourceType.BATCH_FILE + + push_fv_spec = FeatureViewSpec(name="user_push_fv") + push_fv_spec.stream_source.CopyFrom(push_ds) + push_fv_spec.batch_source.CopyFrom(push_batch_ds) + push_fv = FeatureView(spec=push_fv_spec) + + # 2. FeatureView with KafkaSource + kafka_ds = DataSource() + kafka_ds.name = "user_kafka_source" + kafka_ds.type = DataSource.SourceType.STREAM_KAFKA + kafka_ds.kafka_options.topic = "user_events" + + kafka_batch_ds = DataSource() + kafka_batch_ds.name = "user_kafka_batch_source" + kafka_batch_ds.type = DataSource.SourceType.BATCH_FILE + + kafka_fv_spec = FeatureViewSpec(name="user_kafka_fv") + kafka_fv_spec.stream_source.CopyFrom(kafka_ds) + kafka_fv_spec.batch_source.CopyFrom(kafka_batch_ds) + kafka_fv = FeatureView(spec=kafka_fv_spec) + + # 3. FeatureView with KinesisSource + kinesis_ds = DataSource() + kinesis_ds.name = "user_kinesis_source" + kinesis_ds.type = DataSource.SourceType.STREAM_KINESIS + kinesis_ds.kinesis_options.stream_name = "user_stream" + + kinesis_batch_ds = DataSource() + kinesis_batch_ds.name = "user_kinesis_batch_source" + kinesis_batch_ds.type = DataSource.SourceType.BATCH_FILE + + kinesis_fv_spec = FeatureViewSpec(name="user_kinesis_fv") + kinesis_fv_spec.stream_source.CopyFrom(kinesis_ds) + kinesis_fv_spec.batch_source.CopyFrom(kinesis_batch_ds) + kinesis_fv = FeatureView(spec=kinesis_fv_spec) + + registry.feature_views.extend([push_fv, kafka_fv, kinesis_fv]) + + lineage_generator = RegistryLineageGenerator() + direct_relationships, indirect_relationships = ( + lineage_generator.generate_lineage(registry) + ) + + actual_direct = { + ( + rel.source.type.value, + rel.source.name, + rel.target.type.value, + rel.target.name, + ) + for rel in direct_relationships + } + + expected_direct = { + ("dataSource", "user_push_source", "featureView", "user_push_fv"), + ("dataSource", "user_push_batch_source", "featureView", "user_push_fv"), + ("dataSource", "user_kafka_source", "featureView", "user_kafka_fv"), + ("dataSource", "user_kafka_batch_source", "featureView", "user_kafka_fv"), + ("dataSource", "user_kinesis_source", "featureView", "user_kinesis_fv"), + ( + "dataSource", + "user_kinesis_batch_source", + "featureView", + "user_kinesis_fv", + ), + } + + for exp in expected_direct: + assert exp in actual_direct, ( + f"Expected relationship {exp} not found in {actual_direct}" + ) + + # Ensure no reverse edges (e.g. featureView -> dataSource) exist since no upstream_feature_view_names are specified + reverse_edges = [ + rel + for rel in direct_relationships + if rel.source.type == FeastObjectType.FEATURE_VIEW + and rel.target.type == FeastObjectType.DATA_SOURCE + ] + assert len(reverse_edges) == 0 + + def test_python_sdk_feature_view_stream_sources_lineage(self): + """Test high-level Python SDK FeatureView with PushSource, KafkaSource, and KinesisSource + serializes to protobuf and generates expected lineage relationships. + """ + from datetime import timedelta + + from feast import ( + Entity as FeastEntity, + ) + from feast import ( + FeatureView as FeastFeatureView, + ) + from feast import ( + Field as FeastField, + ) + from feast import ( + FileSource as FeastFileSource, + ) + from feast import ( + KafkaSource as FeastKafkaSource, + ) + from feast import ( + KinesisSource as FeastKinesisSource, + ) + from feast import ( + PushSource as FeastPushSource, + ) + from feast import ( + ValueType, + ) + from feast.data_format import ProtoFormat + from feast.types import Float32, String + + user_entity = FeastEntity( + name="user_id", join_keys=["user_id"], value_type=ValueType.STRING + ) + file_source = FeastFileSource( + name="batch_file_source", path="data/batch.parquet" + ) + + # 1. PushSource without upstream_feature_view_names + push_source = FeastPushSource( + name="sdk_push_source", + batch_source=file_source, + description="Push source without upstream views", + ) + push_fv = FeastFeatureView( + name="sdk_push_fv", + entities=[user_entity], + ttl=timedelta(days=1), + schema=[ + FeastField(name="user_id", dtype=String), + FeastField(name="score", dtype=Float32), + ], + source=push_source, + ) + + # 2. KafkaSource + kafka_source = FeastKafkaSource( + name="sdk_kafka_source", + kafka_bootstrap_servers="localhost:9092", + topic="events_topic", + timestamp_field="event_timestamp", + message_format=ProtoFormat("class_path"), + batch_source=file_source, + ) + kafka_fv = FeastFeatureView( + name="sdk_kafka_fv", + entities=[user_entity], + ttl=timedelta(days=1), + schema=[ + FeastField(name="user_id", dtype=String), + FeastField(name="clicks", dtype=Float32), + ], + source=kafka_source, + ) + + # 3. KinesisSource + kinesis_source = FeastKinesisSource( + name="sdk_kinesis_source", + record_format=ProtoFormat("class_path"), + region="us-west-2", + stream_name="events_stream", + timestamp_field="event_timestamp", + batch_source=file_source, + ) + kinesis_fv = FeastFeatureView( + name="sdk_kinesis_fv", + entities=[user_entity], + ttl=timedelta(days=1), + schema=[ + FeastField(name="user_id", dtype=String), + FeastField(name="views", dtype=Float32), + ], + source=kinesis_source, + ) + + registry = Registry() + registry.feature_views.extend( + [ + push_fv.to_proto(), + kafka_fv.to_proto(), + kinesis_fv.to_proto(), + ] + ) + + lineage_generator = RegistryLineageGenerator() + direct_relationships, _ = lineage_generator.generate_lineage(registry) + + actual_direct = { + ( + rel.source.type.value, + rel.source.name, + rel.target.type.value, + rel.target.name, + ) + for rel in direct_relationships + } + + assert ( + "dataSource", + "sdk_push_source", + "featureView", + "sdk_push_fv", + ) in actual_direct + assert ( + "dataSource", + "batch_file_source", + "featureView", + "sdk_push_fv", + ) in actual_direct + assert ( + "dataSource", + "sdk_kafka_source", + "featureView", + "sdk_kafka_fv", + ) in actual_direct + assert ( + "dataSource", + "batch_file_source", + "featureView", + "sdk_kafka_fv", + ) in actual_direct + assert ( + "dataSource", + "sdk_kinesis_source", + "featureView", + "sdk_kinesis_fv", + ) in actual_direct + assert ( + "dataSource", + "batch_file_source", + "featureView", + "sdk_kinesis_fv", + ) in actual_direct diff --git a/sdk/python/tests/unit/test_data_sources.py b/sdk/python/tests/unit/test_data_sources.py index 1d8a263919d..966c95c9318 100644 --- a/sdk/python/tests/unit/test_data_sources.py +++ b/sdk/python/tests/unit/test_data_sources.py @@ -1,3 +1,4 @@ +import warnings from datetime import timedelta import pytest @@ -48,6 +49,121 @@ def test_push_source_without_batch_source(): assert push_source_unproto.name == "test_push_source" +def test_push_source_with_upstream_feature_view_names(): + class DummyFeatureView: + def __init__(self, name: str): + self.name = name + + fv1 = DummyFeatureView("user_transaction_stats") + fv2 = DummyFeatureView("user_credit_profile") + + # Initialize with objects and strings + push_source = PushSource( + name="risk_calc_pipeline", + batch_source=BigQuerySource(table="test.test"), + upstream_feature_view_names=[fv1, fv2], + description="test push source with lineage", + ) + + assert push_source.upstream_feature_view_names == [ + "user_transaction_stats", + "user_credit_profile", + ] + + # Convert to proto and verify + proto = push_source.to_proto() + assert proto.HasField("push_options") + assert list(proto.push_options.upstream_feature_view_names) == [ + "user_transaction_stats", + "user_credit_profile", + ] + + # Deserialize from proto and verify + unproto = PushSource.from_proto(proto) + assert unproto.name == "risk_calc_pipeline" + assert unproto.upstream_feature_view_names == [ + "user_transaction_stats", + "user_credit_profile", + ] + assert unproto.description == "test push source with lineage" + assert unproto == push_source + + # Verify equality and inequality + push_source_same = PushSource( + name="risk_calc_pipeline", + batch_source=BigQuerySource(table="test.test"), + upstream_feature_view_names=[ + "user_transaction_stats", + "user_credit_profile", + ], + description="test push source with lineage", + ) + assert push_source == push_source_same + + push_source_diff = PushSource( + name="risk_calc_pipeline", + batch_source=BigQuerySource(table="test.test"), + upstream_feature_view_names=["user_transaction_stats"], + description="test push source with lineage", + ) + assert push_source != push_source_diff + + +def test_push_source_upstream_validation(): + # Invalid type + with pytest.raises(TypeError, match="Unsupported type for upstream feature view"): + PushSource( + name="invalid_type", + upstream_feature_view_names=[123], # type: ignore[list-item] + ) + + # Empty string + with pytest.raises(ValueError, match="Upstream feature view name cannot be empty"): + PushSource( + name="empty_string", + upstream_feature_view_names=[""], + ) + + # Object with empty name + class BadFeatureView: + def __init__(self): + self.name = " " + + with pytest.raises(ValueError, match="invalid or empty name"): + PushSource( + name="bad_fv_obj", + upstream_feature_view_names=[BadFeatureView()], + ) + + +def test_validate_data_sources_upstream_warning(): + from feast.feature_store import _validate_data_sources + + class DummyFV: + def __init__(self, name: str): + self.name = name + + ps = PushSource( + name="risk_push", + upstream_feature_view_names=["existing_fv", "missing_fv"], + ) + + # Warning expected for "missing_fv" + with pytest.warns( + UserWarning, + match="PushSource 'risk_push' references upstream feature view 'missing_fv'", + ): + _validate_data_sources([ps], all_feature_views=[DummyFV("existing_fv")]) + + # No warning when all are found + with warnings.catch_warnings(): + warnings.simplefilter("error") + _validate_data_sources( + [ps], + all_feature_views=[DummyFV("existing_fv"), DummyFV("missing_fv")], + ) + + def test_request_source_primitive_type_to_proto(): schema = [ Field(name="f1", dtype=Float32), diff --git a/sdk/python/tests/unit/test_metrics.py b/sdk/python/tests/unit/test_metrics.py index 5f345d81f1e..3e00beee6be 100644 --- a/sdk/python/tests/unit/test_metrics.py +++ b/sdk/python/tests/unit/test_metrics.py @@ -13,8 +13,10 @@ # limitations under the License. import http.client +import os import socket import threading +import time from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch @@ -642,6 +644,67 @@ def test_ready_event_set_on_successful_bind(self): mock_event.set.assert_called_once() +@pytest.mark.skipif(not hasattr(os, "fork"), reason="requires os.fork") +@pytest.mark.filterwarnings( + "ignore:This process .* is multi-threaded:DeprecationWarning" +) +class TestMetricsProcessSurvivesForkedChildExit: + """A forked child (e.g. a Gunicorn worker) must not terminate the metrics server on exit.""" + + def test_forked_child_does_not_track_metrics_process(self): + import multiprocessing + + import feast.metrics as m + + # Stand-in for the metrics-server process: a real daemon process that + # starts instantly on every platform, registered by start_metrics_server. + proc = multiprocessing.get_context("fork").Process( + target=time.sleep, args=(60,), daemon=True + ) + mock_event = MagicMock() + mock_event.wait.return_value = True + + with ( + patch("feast.metrics.multiprocessing.Process", return_value=proc), + patch("feast.metrics.multiprocessing.Event", return_value=mock_event), + ): + m.start_metrics_server( + MagicMock(), + start_resource_monitoring=False, + start_freshness_monitoring=False, + ) + try: + pid = os.fork() + if pid == 0: + # A worker's interpreter exit terminates every daemon process in + # active_children(); report whether the metrics process is in it. + tracked = proc in multiprocessing.active_children() + os._exit(1 if tracked else 0) + _, status = os.waitpid(pid, 0) + + assert os.waitstatus_to_exitcode(status) == 0 + assert proc.is_alive() + finally: + proc.kill() + proc.join(timeout=5) + m._metrics_process = None + + +def test_forget_metrics_process_in_child_removes_it_from_registry(): + import multiprocessing + + import feast.metrics as m + + proc = MagicMock() + children = {proc} + with ( + patch.object(multiprocessing.process, "_children", children), + patch.object(m, "_metrics_process", proc), + ): + m._forget_metrics_process_in_child() + assert proc not in children + + class TestMetricsYamlConfig: """Verify metrics config in feature_store.yaml is respected. diff --git a/ui/src/components/RegistryVisualization.test.tsx b/ui/src/components/RegistryVisualization.test.tsx index 01da8f415a3..3dda909b4b8 100644 --- a/ui/src/components/RegistryVisualization.test.tsx +++ b/ui/src/components/RegistryVisualization.test.tsx @@ -4,7 +4,9 @@ import { waitFor } from "@testing-library/react"; import RegistryVisualization from "./RegistryVisualization"; import { feast } from "../protos"; import { FEAST_FCO_TYPES } from "../parsers/types"; -import { EntityRelation } from "../parsers/parseEntityRelationships"; +import parseEntityRelationships, { + EntityRelation, +} from "../parsers/parseEntityRelationships"; import { ThemeProvider } from "../contexts/ThemeContext"; // ReactFlow requires ResizeObserver @@ -370,3 +372,238 @@ describe("RegistryVisualization legend", () => { expect(screen.getByText("vN")).toBeInTheDocument(); }); }); + +describe("parseEntityRelationships PushSource lineage", () => { + test("parses upstreamFeatureViewNames on PushSource into EntityRelation links", () => { + const registry = makeRegistry({ + featureViews: [ + feast.core.FeatureView.create({ + spec: { + name: "user_risk_target_fv", + streamSource: feast.core.DataSource.create({ + name: "risk_calc_pipeline", + pushOptions: feast.core.DataSource.PushOptions.create({ + upstreamFeatureViewNames: [ + "user_transaction_stats", + "user_credit_profile", + ], + }), + }), + }, + }), + ], + dataSources: [ + feast.core.DataSource.create({ + name: "risk_calc_pipeline", + pushOptions: feast.core.DataSource.PushOptions.create({ + upstreamFeatureViewNames: [ + "user_transaction_stats", + "user_credit_profile", + ], + }), + }), + ], + }); + + const links = parseEntityRelationships(registry); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.featureView, + name: "user_transaction_stats", + }, + target: { + type: FEAST_FCO_TYPES.dataSource, + name: "risk_calc_pipeline", + }, + }); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.featureView, + name: "user_credit_profile", + }, + target: { + type: FEAST_FCO_TYPES.dataSource, + name: "risk_calc_pipeline", + }, + }); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.dataSource, + name: "risk_calc_pipeline", + }, + target: { + type: FEAST_FCO_TYPES.featureView, + name: "user_risk_target_fv", + }, + }); + }); + + test("parses streamSource and batchSource on FeatureView with PushSource without upstreamFeatureViewNames", () => { + const registry = makeRegistry({ + featureViews: [ + feast.core.FeatureView.create({ + spec: { + name: "user_push_fv", + streamSource: feast.core.DataSource.create({ + name: "user_push_source", + }), + batchSource: feast.core.DataSource.create({ + name: "user_push_batch_source", + }), + }, + }), + ], + dataSources: [ + feast.core.DataSource.create({ + name: "user_push_source", + }), + feast.core.DataSource.create({ + name: "user_push_batch_source", + }), + ], + }); + + const links = parseEntityRelationships(registry); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.dataSource, + name: "user_push_source", + }, + target: { + type: FEAST_FCO_TYPES.featureView, + name: "user_push_fv", + }, + }); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.dataSource, + name: "user_push_batch_source", + }, + target: { + type: FEAST_FCO_TYPES.featureView, + name: "user_push_fv", + }, + }); + + // Ensure no upstream featureView -> dataSource links exist + const upstreamLinks = links.filter( + (l) => + l.source.type === FEAST_FCO_TYPES.featureView && + l.target.type === FEAST_FCO_TYPES.dataSource, + ); + expect(upstreamLinks).toHaveLength(0); + }); + + test("parses streamSource and batchSource on FeatureView with KafkaSource", () => { + const registry = makeRegistry({ + featureViews: [ + feast.core.FeatureView.create({ + spec: { + name: "user_kafka_fv", + streamSource: feast.core.DataSource.create({ + name: "user_kafka_source", + kafkaOptions: feast.core.DataSource.KafkaOptions.create({ + topic: "user_events", + kafkaBootstrapServers: "localhost:9092", + }), + }), + batchSource: feast.core.DataSource.create({ + name: "user_kafka_batch_source", + }), + }, + }), + ], + dataSources: [ + feast.core.DataSource.create({ + name: "user_kafka_source", + }), + feast.core.DataSource.create({ + name: "user_kafka_batch_source", + }), + ], + }); + + const links = parseEntityRelationships(registry); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.dataSource, + name: "user_kafka_source", + }, + target: { + type: FEAST_FCO_TYPES.featureView, + name: "user_kafka_fv", + }, + }); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.dataSource, + name: "user_kafka_batch_source", + }, + target: { + type: FEAST_FCO_TYPES.featureView, + name: "user_kafka_fv", + }, + }); + }); + + test("parses streamSource and batchSource on FeatureView with KinesisSource", () => { + const registry = makeRegistry({ + featureViews: [ + feast.core.FeatureView.create({ + spec: { + name: "user_kinesis_fv", + streamSource: feast.core.DataSource.create({ + name: "user_kinesis_source", + kinesisOptions: feast.core.DataSource.KinesisOptions.create({ + streamName: "user_stream", + region: "us-west-2", + }), + }), + batchSource: feast.core.DataSource.create({ + name: "user_kinesis_batch_source", + }), + }, + }), + ], + dataSources: [ + feast.core.DataSource.create({ + name: "user_kinesis_source", + }), + feast.core.DataSource.create({ + name: "user_kinesis_batch_source", + }), + ], + }); + + const links = parseEntityRelationships(registry); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.dataSource, + name: "user_kinesis_source", + }, + target: { + type: FEAST_FCO_TYPES.featureView, + name: "user_kinesis_fv", + }, + }); + + expect(links).toContainEqual({ + source: { + type: FEAST_FCO_TYPES.dataSource, + name: "user_kinesis_batch_source", + }, + target: { + type: FEAST_FCO_TYPES.featureView, + name: "user_kinesis_fv", + }, + }); + }); +}); diff --git a/ui/src/components/RegistryVisualization.tsx b/ui/src/components/RegistryVisualization.tsx index 0142943e78d..be265efc4d8 100644 --- a/ui/src/components/RegistryVisualization.tsx +++ b/ui/src/components/RegistryVisualization.tsx @@ -797,6 +797,9 @@ const registryToFlow = ( if (fv.spec?.batchSource?.name) { dataSources.add(fv.spec.batchSource.name); } + if (fv.spec?.streamSource?.name) { + dataSources.add(fv.spec.streamSource.name); + } }); objects.streamFeatureViews?.forEach((sfv) => { @@ -830,6 +833,12 @@ const registryToFlow = ( } }); + // Also include standalone registered data sources + ((objects as any).dataSources || []).forEach((ds: any) => { + const name = ds.spec?.name || ds.name; + if (name) dataSources.add(name); + }); + Array.from(dataSources).forEach((dsName) => { nodes.push({ id: `ds-${dsName}`, diff --git a/ui/src/parsers/parseEntityRelationships.ts b/ui/src/parsers/parseEntityRelationships.ts index 6156b67d895..683e303c29d 100644 --- a/ui/src/parsers/parseEntityRelationships.ts +++ b/ui/src/parsers/parseEntityRelationships.ts @@ -107,6 +107,18 @@ const parseEntityRelationships = (objects: feast.core.Registry) => { }, }); }); + if (fv.spec?.streamSource) { + links.push({ + source: { + type: FEAST_FCO_TYPES["dataSource"], + name: fv.spec.streamSource.name || "", + }, + target: { + type: FEAST_FCO_TYPES["featureView"], + name: fv.spec?.name!, + }, + }); + } if (fv.spec?.batchSource) { links.push({ source: { @@ -244,18 +256,56 @@ const parseEntityRelationships = (objects: feast.core.Registry) => { } }); - // Build data source location index for storage-based matching const allDataSources = [ ...((objects as any).dataSources || []), ...(objects.featureViews || []) - .map((fv: any) => fv.spec?.batchSource) + .flatMap((fv: any) => [fv.spec?.batchSource, fv.spec?.streamSource]) .filter(Boolean), ...(objects.streamFeatureViews || []) .flatMap((sfv: any) => [sfv.spec?.batchSource, sfv.spec?.streamSource]) .filter(Boolean), + ...(((objects as any).labelViews || []) as any[]) + .flatMap((lv: any) => [lv.spec?.source, lv.spec?.batchSource]) + .filter(Boolean), ]; + + // Build data source location index for storage-based matching + // stream sources that don't specify any storage locations are safely skipped const dsLocationIndex = buildDataSourceLocationIndex(allDataSources); + // Upstream FeatureView -> DataSource (PushSource) relationships + // Note: BatchSource and non-push StreamSource do not declare upstreamFeatureViewNames. + const seenPushEdges = new Set(); + allDataSources.forEach((ds: any) => { + const dsObj = ds.spec || ds; + const dsName = dsObj.name; + const pushOpts = dsObj.pushOptions || dsObj.push_options; + const upstreamFvs = + pushOpts?.upstreamFeatureViewNames || + pushOpts?.upstream_feature_view_names; + if (dsName && Array.isArray(upstreamFvs)) { + upstreamFvs.forEach((upstreamFvName: string) => { + const edgeKey = `${upstreamFvName}->${dsName}`; + if (!seenPushEdges.has(edgeKey)) { + seenPushEdges.add(edgeKey); + const isLabelView = labelViewNames.has(upstreamFvName); + links.push({ + source: { + type: isLabelView + ? FEAST_FCO_TYPES["labelView"] + : FEAST_FCO_TYPES["featureView"], + name: upstreamFvName, + }, + target: { + type: FEAST_FCO_TYPES["dataSource"], + name: dsName, + }, + }); + } + }); + } + }); + (objects as any).savedDatasets?.forEach((sd: any) => { if (sd.spec?.featureServiceName) { links.push({