Skip to content

Commit

Permalink
feat: Add to_remote_storage method to RetrievalJob (#2916)
Browse files Browse the repository at this point in the history
* feat: Add to_remote_storage method to RetrievalJob

Signed-off-by: Achal Shah <achals@gmail.com>

* docstring

Signed-off-by: Achal Shah <achals@gmail.com>
  • Loading branch information
achals committed Jul 6, 2022
1 parent dcd8ec9 commit 109ee9c
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 2 deletions.
20 changes: 20 additions & 0 deletions sdk/python/feast/infra/offline_stores/offline_store.py
Expand Up @@ -163,6 +163,26 @@ def metadata(self) -> Optional[RetrievalMetadata]:
"""
pass

def supports_remote_storage_export(self) -> bool:
"""
This method should return True if the RetrievalJob supports `to_remote_storage()`.
"""
return False

def to_remote_storage(self) -> List[str]:
"""
This method should export the result of this RetrievalJob to
remote storage (such as S3, GCS, HDFS, etc).
Implementations of this method should export the results as
multiple parquet files, each file sized appropriately
depending on how much data is being returned by the retrieval
job.
Returns:
A list of parquet file paths in remote storage.
"""
raise NotImplementedError()


class OfflineStore(ABC):
"""
Expand Down
Expand Up @@ -280,7 +280,7 @@ def get_expected_training_df(

@pytest.mark.integration
@pytest.mark.universal_offline_stores
@pytest.mark.parametrize("full_feature_names", [True, False], ids=lambda v: str(v))
@pytest.mark.parametrize("full_feature_names", [True, False], ids=lambda v: f"full:{v}")
def test_historical_features(environment, universal_data_sources, full_feature_names):
store = environment.feature_store

Expand Down Expand Up @@ -410,7 +410,7 @@ def test_historical_features(environment, universal_data_sources, full_feature_n


@pytest.mark.integration
@pytest.mark.universal
@pytest.mark.universal_offline_stores
@pytest.mark.parametrize("full_feature_names", [True, False], ids=lambda v: str(v))
def test_historical_features_with_shared_batch_source(
environment, universal_data_sources, full_feature_names
Expand Down

0 comments on commit 109ee9c

Please sign in to comment.