Adding a new offline store

Overview
Feast makes adding support for a new offline store (database) easy. Developers can simply implement the OfflineStore interface to add support for a new store (other than the existing stores like Parquet files, Redshift, and Bigquery).
In this guide, we will show you how to extend the existing File offline store and use in a feature repo. While we will be implementing a specific store, this guide should be representative for adding support for any new offline store.
The full working code for this guide can be found at feast-dev/feast-custom-offline-store-demo.
The process for using a custom offline store consists of 4 steps:
1.Defining an OfflineStore class.
2.Defining an OfflineStoreConfig class.
3.Defining a RetrievalJob class for this offline store.
4.Defining a DataSource class for the offline store
5.Referencing the OfflineStore in a feature repo's feature_store.yaml file.
6.Testing the OfflineStore class.
1. Defining an OfflineStore class
OfflineStore class names must end with the OfflineStore suffix!
The OfflineStore class contains a couple of methods to read features from the offline store. Unlike the OnlineStore class, Feast does not manage any infrastructure for the offline store.
There are two methods that deal with reading data from the offline storesget_historical_featuresand pull_latest_from_table_or_query.
pull_latest_from_table_or_query is invoked when running materialization (using the feast materialize or feast materialize-incremental commands, or the corresponding FeatureStore.materialize() method. This method pull data from the offline store, and the FeatureStore class takes care of writing this data into the online store.
get_historical_features is invoked when reading values from the offline store using the FeatureStore.get_historical_features() method. Typically, this method is used to retrieve features when training ML models.
pull_all_from_table_or_query is a method that pulls all the data from an offline store from a specified start date to a specified end date.
feast_custom_offline_store/file.py
1
    def get_historical_features(self,
2
                                config: RepoConfig,
3
                                feature_views: List[FeatureView],
4
                                feature_refs: List[str],
5
                                entity_df: Union[pd.DataFrame, str],
6
                                registry: Registry, project: str,
7
                                full_feature_names: bool = False) -> RetrievalJob:
8
        print("Getting historical features from my offline store")
9
        return super().get_historical_features(config,
10
                                               feature_views,
11
                                               feature_refs,
12
                                               entity_df,
13
                                               registry,
14
                                               project,
15
                                               full_feature_names)
16
​
17
    def pull_latest_from_table_or_query(self,
18
                                        config: RepoConfig,
19
                                        data_source: DataSource,
20
                                        join_key_columns: List[str],
21
                                        feature_name_columns: List[str],
22
                                        timestamp_field: str,
23
                                        created_timestamp_column: Optional[str],
24
                                        start_date: datetime,
25
                                        end_date: datetime) -> RetrievalJob:
26
        print("Pulling latest features from my offline store")
27
        return super().pull_latest_from_table_or_query(config,
28
                                                       data_source,
29
                                                       join_key_columns,
30
                                                       feature_name_columns,
31
                                                       timestamp_field=timestamp_field,
32
                                                       created_timestamp_column,
33
                                                       start_date,
34
                                                       end_date)
Copied!
2. Defining an OfflineStoreConfig class
Additional configuration may be needed to allow the OfflineStore to talk to the backing store. For example, Redshift needs configuration information like the connection information for the Redshift instance, credentials for connecting to the database, etc.
To facilitate configuration, all OfflineStore implementations are required to also define a corresponding OfflineStoreConfig class in the same file. This OfflineStoreConfig class should inherit from the FeastConfigBaseModel class, which is defined here.
The FeastConfigBaseModel is a pydantic class, which parses yaml configuration into python objects. Pydantic also allows the model classes to define validators for the config classes, to make sure that the config classes are correctly defined.
This config class must container a type field, which contains the fully qualified class name of its corresponding OfflineStore class.
Additionally, the name of the config class must be the same as the OfflineStore class, with the Config suffix.
An example of the config class for the custom file offline store :
feast_custom_offline_store/file.py
1
class CustomFileOfflineStoreConfig(FeastConfigBaseModel):
2
    """ Custom offline store config for local (file-based) store """
3
​
4
    type: Literal["feast_custom_offline_store.file.CustomFileOfflineStore"] \
5
        = "feast_custom_offline_store.file.CustomFileOfflineStore"
Copied!
This configuration can be specified in the feature_store.yaml as follows:
feature_repo/feature_store.yaml
1
type: feast_custom_offline_store.file.CustomFileOfflineStore
Copied!
This configuration information is available to the methods of the OfflineStore, via theconfig: RepoConfig parameter which is passed into the methods of the OfflineStore interface, specifically at the config.offline_store field of the config parameter.
feast_custom_offline_store/file.py
1
    def get_historical_features(self,
2
                                config: RepoConfig,
3
                                feature_views: List[FeatureView],
4
                                feature_refs: List[str],
5
                                entity_df: Union[pd.DataFrame, str],
6
                                registry: Registry, project: str,
7
                                full_feature_names: bool = False) -> RetrievalJob:
8
​
9
        offline_store_config = config.offline_store
10
        assert isinstance(offline_store_config, CustomFileOfflineStoreConfig)
11
        store_type = offline_store_config.type
Copied!
3. Defining a RetrievalJob class
The offline store methods aren't expected to perform their read operations eagerly. Instead, they are expected to execute lazily, and they do so by returning a RetrievalJob instance, which represents the execution of the actual query against the underlying store.
Custom offline stores may need to implement their own instances of the RetrievalJob interface.
The RetrievalJob interface exposes two methods - to_df and to_arrow. The expectation is for the retrieval job to be able to return the rows read from the offline store as a parquet DataFrame, or as an Arrow table respectively.
feast_custom_offline_store/file.py
1
class CustomFileRetrievalJob(RetrievalJob):
2
    def __init__(self, evaluation_function: Callable):
3
        """Initialize a lazy historical retrieval job"""
4
​
5
        # The evaluation function executes a stored procedure to compute a historical retrieval.
6
        self.evaluation_function = evaluation_function
7
​
8
    def to_df(self):
9
        # Only execute the evaluation function to build the final historical retrieval dataframe at the last moment.
10
        print("Getting a pandas DataFrame from a File is easy!")
11
        df = self.evaluation_function()
12
        return df
13
​
14
    def to_arrow(self):
15
        # Only execute the evaluation function to build the final historical retrieval dataframe at the last moment.
16
        print("Getting a pandas DataFrame from a File is easy!")
17
        df = self.evaluation_function()
18
        return pyarrow.Table.from_pandas(df)
Copied!
4. Defining a DataSource class for the offline store
Before this offline store can be used as the batch source for a feature view in a feature repo, a subclass of the DataSource base class needs to be defined. This class is responsible for holding information needed by specific feature views to support reading historical values from the offline store. For example, a feature view using Redshift as the offline store may need to know which table contains historical feature values.
The data source class should implement two methods - from_proto, and to_proto.
For custom offline stores that are not being implemented in the main feature repo, the custom_options field should be used to store any configuration needed by the data source. In this case, the implementer is responsible for serializing this configuration into bytes in the to_proto method and reading the value back from bytes in the from_proto method.
feast_custom_offline_store/file.py
1
class CustomFileDataSource(FileSource):
2
    """Custom data source class for local files"""
3
    def __init__(
4
        self,
5
        timestamp_field: Optional[str] = "",
6
        path: Optional[str] = None,
7
        field_mapping: Optional[Dict[str, str]] = None,
8
        created_timestamp_column: Optional[str] = "",
9
        date_partition_column: Optional[str] = "",
10
    ):
11
        super(CustomFileDataSource, self).__init__(
12
            timestamp_field=timestamp_field,
13
            created_timestamp_column,
14
            field_mapping,
15
            date_partition_column,
16
        )
17
        self._path = path
18
​
19
​
20
    @staticmethod
21
    def from_proto(data_source: DataSourceProto):
22
        custom_source_options = str(
23
            data_source.custom_options.configuration, encoding="utf8"
24
        )
25
        path = json.loads(custom_source_options)["path"]
26
        return CustomFileDataSource(
27
            field_mapping=dict(data_source.field_mapping),
28
            path=path,
29
            timestamp_field=data_source.timestamp_field,
30
            created_timestamp_column=data_source.created_timestamp_column,
31
            date_partition_column=data_source.date_partition_column,
32
        )
33
​
34
    def to_proto(self) -> DataSourceProto:
35
        config_json = json.dumps({"path": self.path})
36
        data_source_proto = DataSourceProto(
37
            type=DataSourceProto.CUSTOM_SOURCE,
38
            custom_options=DataSourceProto.CustomSourceOptions(
39
                configuration=bytes(config_json, encoding="utf8")
40
            ),
41
        )
42
​
43
        data_source_proto.timestamp_field = self.timestamp_field
44
        data_source_proto.created_timestamp_column = self.created_timestamp_column
45
        data_source_proto.date_partition_column = self.date_partition_column
46
​
47
        return data_source_proto
Copied!
5. Using the custom offline store
After implementing these classes, the custom offline store can be used by referencing it in a feature repo's feature_store.yaml file, specifically in the offline_store field. The value specified should be the fully qualified class name of the OfflineStore.
As long as your OfflineStore class is available in your Python environment, it will be imported by Feast dynamically at runtime.
To use our custom file offline store, we can use the following feature_store.yaml:
feature_repo/feature_store.yaml
1
project: test_custom
2
registry: data/registry.db
3
provider: local
4
offline_store:
5
    type: feast_custom_offline_store.file.CustomFileOfflineStore
Copied!
If additional configuration for the offline store is **not **required, then we can omit the other fields and only specify the type of the offline store class as the value for the offline_store.
feature_repo/feature_store.yaml
1
project: test_custom
2
registry: data/registry.db
3
provider: local
4
offline_store: feast_custom_offline_store.file.CustomFileOfflineStore
Copied!
Finally, the custom data source class can be use in the feature repo to define a data source, and refer to in a feature view definition.
feature_repo/repo.py
1
pdriver_hourly_stats = CustomFileDataSource(
2
    path="feature_repo/data/driver_stats.parquet",
3
    timestamp_field="event_timestamp",
4
    created_timestamp_column="created",
5
)
6
​
7
​
8
driver_hourly_stats_view = FeatureView(
9
    source=driver_hourly_stats,
10
    ...
11
)
Copied!
6. Testing the OfflineStore class
Even if you have created the OfflineStore class in a separate repo, you can still test your implementation against the Feast test suite, as long as you have Feast as a submodule in your repo. In the Feast submodule, we can run all the unit tests with:
1
make test
Copied!
The universal tests, which are integration tests specifically intended to test offline and online stores, can be run with:
1
make test-python-universal
Copied!
The unit tests should succeed, but the universal tests will likely fail. The tests are parametrized based on the FULL_REPO_CONFIGS variable defined in sdk/python/tests/integration/feature_repos/repo_configuration.py. To overwrite these configurations, you can simply create your own file that contains a FULL_REPO_CONFIGS, and point Feast to that file by setting the environment variable FULL_REPO_CONFIGS_MODULE to point to that file. The main challenge there will be to write a DataSourceCreator for the offline store. In this repo, the file that overwrites FULL_REPO_CONFIGS is feast_custom_offline_store/feast_tests.py, so you would run
1
export FULL_REPO_CONFIGS_MODULE='feast_custom_offline_store.feast_tests'
2
make test-python-universal
Copied!
to test the offline store against the Feast universal tests. You should notice that some of the tests actually fail; this indicates that there is a mistake in the implementation of this offline store!
How-to Guides - Previous
Adding a new online store
Next - How-to Guides
Adding or reusing tests
Last modified 2mo ago
Export as PDF
Copy link
Edit on GitHub
Contents
Overview
1. Defining an OfflineStore class
2. Defining an OfflineStoreConfig class
3. Defining a RetrievalJob class
4. Defining a DataSource class for the offline store
5. Using the custom offline store
6. Testing the OfflineStore class