apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store module

class apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.VertexAIFeatureStoreEnrichmentHandler(project: str, location: str, api_endpoint: str, feature_store_name: str, feature_view_name: str, row_key: str, *, exception_level: ExceptionLevel = ExceptionLevel.WARN, **kwargs)[source]

Bases: EnrichmentSourceHandler[Row, Row]

Enrichment handler to interact with Vertex AI Feature Store.

Use this handler with apache_beam.transforms.enrichment.Enrichment transform when the Vertex AI Feature Store is set up for Bigtable Online serving.

With the Bigtable Online serving approach, the client fetches all the available features for an entity-id. The entity-id is extracted from the row_key field in the input beam.Row object. To filter the features to enrich, use the join_fn param in apache_beam.transforms.enrichment.Enrichment.

NOTE: The default severity to report exceptions is logging a warning. For

this handler, Vertex AI client returns the same exception Requested entity was not found even though the feature store doesn’t exist. So make sure the feature store instance exists or set exception_level as ExceptionLevel.RAISE.

Initializes an instance of VertexAIFeatureStoreEnrichmentHandler.

Parameters:

  • project (str) – The GCP project-id for the Vertex AI Feature Store.

  • location (str) – The region for the Vertex AI Feature Store.

  • api_endpoint (str) – The API endpoint for the Vertex AI Feature Store.

  • feature_store_name (str) – The name of the Vertex AI Feature Store.

  • feature_view_name (str) – The name of the feature view within the Feature Store.

  • row_key (str) – The row key field name containing the unique id for the feature values.

  • exception_level – a enum.Enum value from apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel to set the level when an empty row is returned from the BigTable query. Defaults to ExceptionLevel.WARN.

  • kwargs – Optional keyword arguments to configure the aiplatform.gapic.FeatureOnlineStoreServiceClient.

get_cache_key(request: Row) → str[source]

Returns a string formatted with unique entity-id for the feature values.

class apache_beam.transforms.enrichment_handlers.vertex_ai_feature_store.VertexAIFeatureStoreLegacyEnrichmentHandler(project: str, location: str, api_endpoint: str, feature_store_id: str, entity_type_id: str, feature_ids: list[str], row_key: str, *, exception_level: ExceptionLevel = ExceptionLevel.WARN, **kwargs)[source]

Bases: EnrichmentSourceHandler

Enrichment handler to interact with Vertex AI Feature Store (Legacy).

Use this handler with apache_beam.transforms.enrichment.Enrichment transform for the Vertex AI Feature Store (Legacy).

By default, it fetches all the features values for an entity-id. The entity-id is extracted from the row_key field in the input beam.Row object.You can specify the features names using feature_ids to fetch specific features.

Initializes an instance of VertexAIFeatureStoreLegacyEnrichmentHandler.

Parameters:

  • project (str) – The GCP project for the Vertex AI Feature Store (Legacy).

  • location (str) – The region for the Vertex AI Feature Store (Legacy).

  • api_endpoint (str) – The API endpoint for the Vertex AI Feature Store (Legacy).

  • feature_store_id (str) – The id of the Vertex AI Feature Store (Legacy).

  • entity_type_id (str) – The entity type of the feature store.

  • feature_ids (list[str]) – A list of feature-ids to fetch from the Feature Store.

  • row_key (str) – The row key field name containing the entity id for the feature values.

  • exception_level – a enum.Enum value from apache_beam.transforms.enrichment_handlers.utils.ExceptionLevel to set the level when an empty row is returned from the BigTable query. Defaults to ExceptionLevel.WARN.

  • kwargs – Optional keyword arguments to configure the aiplatform.gapic.FeaturestoreOnlineServingServiceClient.

get_cache_key(request: Row) → str[source]

Returns a string formatted with unique entity-id for the feature values.