feast-dev · woop · Aug 22, 2021 · Aug 22, 2021 · Aug 22, 2021 · Aug 22, 2021
diff --git a/.github/workflows/test_provider.yml b/.github/workflows/test_provider.yml
@@ -0,0 +1,25 @@
+name: test-provider
+
+on: [push, pull_request]
+jobs:
+  test-provider:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [ 3.7, 3.8, 3.9 ]
+        os: [ ubuntu-latest, macOS-latest]
+    env:
+      OS: ${{ matrix.os }}
+      PYTHON: ${{ matrix.python-version }}
+    steps:
+      - uses: actions/checkout@v2
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+          architecture: x64
+      - name: Install dependencies
+        run: pip install -r requirements.txt
+      - name: Run tests
+        run: pytest test_custom_provider.py
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,5 @@
+.idea
+*.pyc
+**__pycache__
+basic_feature_repo/online_store.db
+basic_feature_repo/registry.db
diff --git a/README.md b/README.md
@@ -0,0 +1,56 @@
+# Feast Custom Provider
+[![test-provider](https://github.com/feast-dev/feast-custom-provider-demo/actions/workflows/test_provider.yml/badge.svg?branch=master)](https://github.com/feast-dev/feast-custom-provider-demo/actions/workflows/test_provider.yml)
+
+### Overview
+
+This repository demonstrates how developers can create their own custom `providers` for Feast. Custom providers can be
+used like plugins which allow Feast users to execute any custom logic. Typical examples include
+* Launching custom streaming ingestion jobs (Spark, Beam)
+* Launching custom batch ingestion (materialization) jobs (Spark, Beam)
+* Adding custom validation to feature repositories during `feast apply`
+* Adding custom infrastructure setup logic which runs during `feast apply`
+* Extending Feast commands with in-house metrics, logging, or tracing
+
+### Why create a custom provider?
+
+All Feast operations execute through a provider. Operations like materializing data from the offline to the online
+store, updating infrastructure like databases, launching streaming ingestion jobs, building training datasets, and
+reading features from the online store.
+
+Feast comes with providers built in, e.g, LocalProvider, GcpProvider, and AwsProvider. However, users can develop their
+own providers by creating a class that implements the contract in the [Provider class](https://github.com/feast-dev/feast/blob/745a1b43d20c0169b675b1f28039854205fb8180/sdk/python/feast/infra/provider.py#L22).
+
+Most developers, however, simply want to add new logic to Feast and don't necessarily want to create a whole provider on
+their own. The fastest way to add custom logic to Feast is to extend an existing provider. The most generic
+provider is the LocalProvider, which contains no custom logic specific to a cloud environment.
+
+### What is included in this repository?
+
+* [feast_custom_provider/](feast_custom_provider): An example of a custom provider, `MyCustomProvider`, which extends the Feast
+`LocalProvider`. This example provider simply prints messages to the console.
+* [basic_feature_repo/](basic_feature_repo): A simple feature repository that is used to test the custom provider. The repository has been configured to use the custom provider as part of it's `feature_store.yaml`
+* [test_custom_provider.py](test_custom_provider.py): A test case that uses `MyCustomProvider` through the `basic_feature_repo/`
+
+### Testing the custom provider in this repository
+
+Run the following commands to test the custom provider ([MyCustomProvider](https://github.com/feast-dev/feast-custom-provider-demo/blob/master/feast_custom_provider/custom_provider.py))
+
+```bash
+pip install -r requirements.txt
+```
+
+```
+pytest test_custom_provider.py
+```
+
+It is also possible to run Feast CLI command, which in turn will call the provider. It may be necessary to add the 
+`PYTHONPATH` to the path where your provider module is stored.
+```
+PYTHONPATH=$PYTHONPATH:/$(pwd) feast -c basic_feature_repo apply
+```
+```
+Registered entity driver_id
+Registered feature view driver_hourly_stats
+Deploying infrastructure for driver_hourly_stats
+Launching custom streaming jobs is pretty easy...
+```
diff --git a/basic_feature_repo/data/driver_stats.parquet b/basic_feature_repo/data/driver_stats.parquet
diff --git a/basic_feature_repo/feature_store.yaml b/basic_feature_repo/feature_store.yaml
@@ -0,0 +1,8 @@
+project: repo
+registry: registry.db
+provider: feast_custom_provider.custom_provider.MyCustomProvider
+online_store:
+    type: sqlite
+    path: online_store.db
+offline_store:
+    type: file
diff --git a/basic_feature_repo/repo.py b/basic_feature_repo/repo.py
@@ -0,0 +1,24 @@
+from datetime import timedelta
+
+from feast import Entity, Feature, FeatureView, FileSource, ValueType
+
+driver_hourly_stats = FileSource(
+    path="basic_feature_repo/data/driver_stats.parquet",
+    event_timestamp_column="event_timestamp",
+    created_timestamp_column="created",
+)
+
+driver = Entity(name="driver_id", value_type=ValueType.INT64, description="driver id",)
+
+driver_hourly_stats_view = FeatureView(
+    name="driver_hourly_stats",
+    entities=["driver_id"],
+    ttl=timedelta(days=365),
+    features=[
+        Feature(name="conv_rate", dtype=ValueType.FLOAT),
+        Feature(name="acc_rate", dtype=ValueType.FLOAT),
+        Feature(name="avg_daily_trips", dtype=ValueType.INT64),
+    ],
+    online=True,
+    batch_source=driver_hourly_stats,
+)
diff --git a/feast_custom_provider/README.md b/feast_custom_provider/README.md
@@ -0,0 +1,3 @@
+# Feast Custom Provider
+
+This custom Feast provider extends the Feast LocalProvider by printing out logs during common Feast operations.
diff --git a/feast_custom_provider/custom_provider.py b/feast_custom_provider/custom_provider.py
@@ -0,0 +1,121 @@
+from datetime import datetime
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+import pandas
+from feast.entity import Entity
+from feast.feature_table import FeatureTable
+from feast.feature_view import FeatureView
+from feast.infra.local import LocalProvider
+from feast.infra.offline_stores.offline_store import RetrievalJob
+from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
+from feast.protos.feast.types.Value_pb2 import Value as ValueProto
+from feast.registry import Registry
+from feast.repo_config import RepoConfig
+from tqdm import tqdm
+
+
+class MyCustomProvider(LocalProvider):
+    def __init__(self, config: RepoConfig, repo_path):
+        super().__init__(config)
+        # Add your custom init code here. This code runs on every feast operation.
+
+    def update_infra(
+        self,
+        project: str,
+        tables_to_delete: Sequence[Union[FeatureTable, FeatureView]],
+        tables_to_keep: Sequence[Union[FeatureTable, FeatureView]],
+        entities_to_delete: Sequence[Entity],
+        entities_to_keep: Sequence[Entity],
+        partial: bool,
+    ):
+        # The update_infra method will be run during "feast apply" and is used to set up databases or launch
+        # long-running jobs on a per table/view basis. This method should also clean up infrastructure that is unused
+        # when feature views or tables are deleted. Examples of operations that update_infra typically fulfills
+        # * Creating, updating, or removing database schemas for tables in an online store
+        # * Launching a streaming ingestion job that writes features into an online store
+
+        # Replace the code below in order to define your own custom infrastructure update operations
+        super().update_infra(
+            project,
+            tables_to_delete,
+            tables_to_keep,
+            entities_to_delete,
+            entities_to_keep,
+            partial,
+        )
+        print("Launching custom streaming jobs is pretty easy...")
+
+    def teardown_infra(
+        self,
+        project: str,
+        tables: Sequence[Union[FeatureTable, FeatureView]],
+        entities: Sequence[Entity],
+    ):
+        # teardown_infra should remove all deployed infrastructure
+
+        # Replace the code below in order to define your own custom teardown operations
+        super().teardown_infra(project, tables, entities)
+
+    def online_write_batch(
+        self,
+        config: RepoConfig,
+        table: Union[FeatureTable, FeatureView],
+        data: List[
+            Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
+        ],
+        progress: Optional[Callable[[int], Any]],
+    ) -> None:
+        # online_write_batch writes feature values to the online store
+        super().online_write_batch(config, table, data, progress)
+
+    def materialize_single_feature_view(
+        self,
+        config: RepoConfig,
+        feature_view: FeatureView,
+        start_date: datetime,
+        end_date: datetime,
+        registry: Registry,
+        project: str,
+        tqdm_builder: Callable[[int], tqdm],
+    ) -> None:
+        # materialize_single_feature_view loads the latest feature values for a specific feature value from the offline
+        # store into the online store.
+        # This method can be overridden to also launch custom batch ingestion jobs that loads the latest batch feature
+        # values into the online store.
+
+        # Replace the line below with your custom logic in order to launch your own batch ingestion job
+        super().materialize_single_feature_view(
+            config, feature_view, start_date, end_date, registry, project, tqdm_builder
+        )
+        print("Launching custom batch jobs is pretty easy...")
+
+    def get_historical_features(
+        self,
+        config: RepoConfig,
+        feature_views: List[FeatureView],
+        feature_refs: List[str],
+        entity_df: Union[pandas.DataFrame, str],
+        registry: Registry,
+        project: str,
+        full_feature_names: bool,
+    ) -> RetrievalJob:
+        # get_historical_features returns a training dataframe from the offline store
+        return super().get_historical_features(
+            config,
+            feature_views,
+            feature_refs,
+            entity_df,
+            registry,
+            project,
+            full_feature_names,
+        )
+
+    def online_read(
+        self,
+        config: RepoConfig,
+        table: Union[FeatureTable, FeatureView],
+        entity_keys: List[EntityKeyProto],
+        requested_features: List[str] = None,
+    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
+        # get_historical_features returns a training dataframe from the offline store
+        return super().online_read(config, table, entity_keys, requested_features)
diff --git a/feast_custom_provider/setup.py b/feast_custom_provider/setup.py
@@ -0,0 +1,9 @@
+from distutils.core import setup
+
+setup(
+    name="feast_custom_provider",
+    version="0.0.1",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    install_requires=["feast==0.12.1"],
+)
diff --git a/requirements.txt b/requirements.txt
@@ -0,0 +1,2 @@
+feast_custom_provider/
+pytest==6.2.4
diff --git a/test_custom_provider.py b/test_custom_provider.py
@@ -0,0 +1,39 @@
+import os
+from datetime import datetime
+
+from feast import FeatureStore
+
+from basic_feature_repo.repo import driver, driver_hourly_stats_view
+
+
+def test_end_to_end():
+    fs = FeatureStore("basic_feature_repo/")
+
+    # apply repository
+    fs.apply([driver, driver_hourly_stats_view])
+
+    # load data into online store
+    fs.materialize_incremental(end_date=datetime.now())
+
+    # Read features from online store
+    feature_vector = fs.get_online_features(
+        features=["driver_hourly_stats:conv_rate"], entity_rows=[{"driver_id": 1001}]
+    ).to_dict()
+    conv_rate = feature_vector["conv_rate"][0]
+    assert conv_rate > 0
+
+    # tear down feature store
+    fs.teardown()
+
+
+def test_cli():
+    os.system(
+        "PYTHONPATH=$PYTHONPATH:/$(pwd) feast -c basic_feature_repo apply > output"
+    )
+    with open("output", "r") as f:
+        output = f.read()
+
+    if "Launching custom streaming jobs is pretty easy" not in output:
+        raise Exception(
+            'Failed to successfully use provider from CLI. See "output" for more details.'
+        )
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,3 @@
		# Feast Custom Provider

		This custom Feast provider extends the Feast LocalProvider by printing out logs during common Feast operations.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		feast_custom_provider/
Comment thread woop marked this conversation as resolved.
		pytest==6.2.4