feast.FeastDataset

kedro_datasets_experimental.feast.FeastDataset

FeastDataset(
    *,
    repo,
    credentials=None,
    load_args=None,
    save_args=None,
    metadata=None
)

Bases: AbstractDataset

Kedro dataset backed by a Feast feature view or feature service.

Load returns a FeastFeatureSource for point-in-time queries against the offline store. Save persists a DataFrame to a feature view's batch source via Feast's write_to_offline_store (which matches the DataFrame to the feature view's schema and appends) and, when write_mode is online_and_offline, also to the online store. Feast infers the write schema from the existing table; set create_table: true to bootstrap a (BigQuery) table from the feature view schema when it does not yet exist.

Examples:

Using the YAML API:

feast_features:
    type: kedro_datasets_experimental.feast.FeastDataset
    repo:  # forwarded to feast.RepoConfig
        registry: gs://bucket/feast
        project: project_name
        provider: gcp # default is local
        offline_store:
            type: bigquery # default is bigquery
            location: EU
        online_store:  # optional; required for online_and_offline writes
            type: sqlite
            path: /tmp/online_store.db
    save_args:
    feature_view_name: features
    write_mode: online_and_offline  # or "offline" (default)
    create_table: true  # bootstrap the BigQuery table from the FV schema
    load_args:
    feature_view_name: features

Parameters:

• repo (dict[str, Any]) –

  Feast repo configuration, forwarded to feast.RepoConfig (which validates it). Any RepoConfig field is accepted, e.g. registry (a path string or dict), project, provider and offline_store. provider defaults to "local" and offline_store to {"type": "bigquery"} when not set; either can be overridden via repo. online_store defaults to disabled (None); set it in repo to enable online_and_offline writes.

• credentials (dict[str, Any] | str | None, default: None ) –

  Not supported; Feast authenticates via Application Default Credentials (the ambient environment). Passing anything other than None raises NotImplementedError.

• load_args (dict[str, Any] | None, default: None ) –

  Must contain either feature_view_name or feature_service_name identifying what load retrieves.

• save_args (dict[str, Any] | None, default: None ) –

  Must contain feature_view_name — the feature view to write to. Optional write_mode selects where to write: "offline" (default) writes only the offline store; "online_and_offline" also pushes to the online store (requires an online_store in repo). Optional create_table (default False): when True, create the backing table from the feature view schema if it does not exist (BigQuery sources only; raises a DatasetError otherwise).

Source code in kedro_datasets_experimental/feast/feast_dataset.py

def __init__(
    self,
    *,
    repo: dict[str, Any],
    credentials: dict[str, Any] | str | None = None,
    load_args: dict[str, Any] | None = None,
    save_args: dict[str, Any] | None = None,
    metadata: dict[str, Any] | None = None,
) -> None:
    """Create a new ``FeastDataset``.

    Args:
        repo: Feast repo configuration, forwarded to ``feast.RepoConfig``
            (which validates it). Any ``RepoConfig`` field is accepted, e.g.
            ``registry`` (a path string or dict), ``project``, ``provider``
            and ``offline_store``. ``provider`` defaults to ``"local"`` and
            ``offline_store`` to ``{"type": "bigquery"}`` when not set;
            either can be overridden via ``repo``. ``online_store`` defaults
            to disabled (``None``); set it in ``repo`` to enable
            ``online_and_offline`` writes.
        credentials: Not supported; Feast authenticates via Application
            Default Credentials (the ambient environment). Passing anything
            other than ``None`` raises ``NotImplementedError``.
        load_args: Must contain either ``feature_view_name`` or
            ``feature_service_name`` identifying what ``load`` retrieves.
        save_args: Must contain ``feature_view_name`` — the feature view to
            write to. Optional ``write_mode`` selects where to write:
            ``"offline"`` (default) writes only the offline store;
            ``"online_and_offline"`` also pushes to the online store
            (requires an ``online_store`` in ``repo``). Optional
            ``create_table`` (default ``False``): when ``True``, create the
            backing table from the feature view schema if it does not exist
            (BigQuery sources only; raises a ``DatasetError`` otherwise).
    """
    if credentials is not None:
        raise NotImplementedError("`FeastDataset` supports only Application Default Credentials")

    repo = repo or {}
    repo_config = RepoConfig(
        **{
            "provider": "local",
            "offline_store": {"type": "bigquery"},
            "online_store": None,
            **repo,
        }
    )

    self._repo = repo
    self._feature_store = FeatureStore(config=repo_config)
    self._load_args = load_args or {}
    self._save_args = save_args or {}
    self.metadata = metadata

_WRITE_MODES class-attribute instance-attribute

_WRITE_MODES = frozenset({'offline', 'online_and_offline'})

_feature_store instance-attribute

_feature_store = FeatureStore(config=repo_config)

_load_args instance-attribute

_load_args = load_args or {}

_repo instance-attribute

_repo = repo

_save_args instance-attribute

_save_args = save_args or {}

metadata instance-attribute

metadata = metadata

_build_bq_schema staticmethod

_build_bq_schema(feature_view, timestamp_field)

Build the BigQuery schema: entity keys, then features, then timestamp.

Source code in kedro_datasets_experimental/feast/feast_dataset.py

@staticmethod
def _build_bq_schema(
    feature_view: FeatureView, timestamp_field: str
) -> list[bigquery.SchemaField]:
    """Build the BigQuery schema: entity keys, then features, then timestamp."""

    def bq_type(dtype: Any) -> str:
        key = dtype.name if hasattr(dtype, "name") else type(dtype).__name__
        return _FEAST_FIELD_TYPE_TO_BQ.get(key, "STRING")

    fields: list[bigquery.SchemaField] = []
    entity_names: set[str] = set()
    for field in feature_view.entity_columns:
        fields.append(bigquery.SchemaField(field.name, bq_type(field.dtype)))
        entity_names.add(field.name)

    for field in feature_view.schema:
        if field.name not in entity_names:
            fields.append(
                bigquery.SchemaField(
                    field.name, bq_type(field.dtype), mode="NULLABLE"
                )
            )

    fields.append(bigquery.SchemaField(timestamp_field, "TIMESTAMP"))
    return fields

_create_bigquery_table

_create_bigquery_table(feature_view_name)

Create the feature view's backing BigQuery table if it doesn't exist.

The table schema is derived from the feature view (entity keys, features, then the event timestamp). Raises aDatasetError for non-BigQuery sources.

Source code in kedro_datasets_experimental/feast/feast_dataset.py

def _create_bigquery_table(self, feature_view_name: str) -> None:
    """Create the feature view's backing BigQuery table if it doesn't exist.

    The table schema is derived from the feature view (entity keys, features,
    then the event timestamp). Raises a``DatasetError`` for non-BigQuery sources.
    """
    feature_view = self._feature_store.get_feature_view(feature_view_name)
    source = feature_view.source
    batch_source = getattr(source, "batch_source", source)
    if not isinstance(batch_source, BigQuerySource):
        msg = (
            f"save_args['create_table'] is only supported for BigQuerySource "
            f"batch sources, got {type(batch_source).__name__!r} for "
            f"'{feature_view_name}'."
        )
        raise DatasetError(msg)
    if not batch_source.table:
        msg = (
            f"Feature view '{feature_view_name}' uses a query-based source "
            f"with no table to create."
        )
        raise DatasetError(msg)

    timestamp_field = batch_source.timestamp_field
    schema = self._build_bq_schema(feature_view, timestamp_field)

    # Run against the project the source table lives in (a fully-qualified
    # `project.dataset.table` ref), falling back to the offline store's.
    offline = self._feature_store.config.offline_store
    parts = batch_source.table.split(".")
    is_fully_qualified = len(parts) == _FULLY_QUALIFIED_BQ_TABLE_PARTS
    table_project = (
        parts[0]
        if is_fully_qualified
        else (offline.billing_project_id or offline.project_id)
    )

    if table_project is None:
        msg = (
            f"Cannot resolve a project for table '{batch_source.table}'; set "
            f"'billing_project_id' or 'project_id' on the offline_store, or use "
            f"a fully-qualified 'project.dataset.table' reference."
        )
        raise DatasetError(msg)

    table_id = (
        batch_source.table
        if is_fully_qualified
        else f"{table_project}.{batch_source.table}"
    )
    table = bigquery.Table(table_id, schema=schema)
    client = bigquery.Client(project=table_project, location=offline.location)
    client.create_table(table, exists_ok=True)

_describe

_describe()

Source code in kedro_datasets_experimental/feast/feast_dataset.py

def _describe(self) -> dict[str, Any]:
    return {
        "repo": self._repo,
        "load_args": self._load_args,
        "save_args": self._save_args,
    }

load

load()

Source code in kedro_datasets_experimental/feast/feast_dataset.py

def load(self) -> FeastFeatureSource:
    if (feature_view_name := self._load_args.get("feature_view_name")) is not None:
        feature_view = self._feature_store.get_feature_view(feature_view_name)
        features = [
            f"{feature_view.name}:{field.name}" for field in feature_view.features
        ]
        return FeastFeatureSource(self._feature_store, features)

    if (
        feature_service_name := self._load_args.get("feature_service_name")
    ) is not None:
        feature_service = self._feature_store.get_feature_service(
            feature_service_name
        )
        return FeastFeatureSource(self._feature_store, feature_service)

    msg = (
        "load_args must contain either a `feature_view_name` or "
        "`feature_service_name`"
    )
    raise DatasetError(msg)

save

save(data)

Source code in kedro_datasets_experimental/feast/feast_dataset.py

def save(self, data: pd.DataFrame) -> None:
    feature_view_name = self._save_args.get("feature_view_name")
    if feature_view_name is None:
        msg = "save_args must contain a `feature_view_name`"
        raise DatasetError(msg)

    write_mode = self._save_args.get("write_mode", "offline")
    if write_mode not in self._WRITE_MODES:
        msg = (
            f"save_args['write_mode'] must be one of "
            f"{sorted(self._WRITE_MODES)}, got {write_mode!r}."
        )
        raise DatasetError(msg)

    # Optionally bootstrap the backing table from the feature view schema.
    # Feast's write infers its schema from the *existing* table, so this
    # must happen first. Only BigQuery-backed sources are supported for now.
    if self._save_args.get("create_table", False):
        self._create_bigquery_table(feature_view_name)

    # Persist the DataFrame into the feature view's batch source: Feast
    # matches/reorders columns to the Feast schema and appends. Fails if
    # columns don't match.
    self._feature_store.write_to_offline_store(
        feature_view_name, data, reorder_columns=True
    )

    # Optionally also push to the online store (requires an online_store to
    # be configured in `repo`; it is disabled by default).
    if write_mode == "online_and_offline":
        self._feature_store.write_to_online_store(feature_view_name, data)