Skip to content

Commit 109ee9c

Browse files
authored
feat: Add to_remote_storage method to RetrievalJob (#2916)
* feat: Add to_remote_storage method to RetrievalJob Signed-off-by: Achal Shah <[email protected]> * docstring Signed-off-by: Achal Shah <[email protected]>
1 parent dcd8ec9 commit 109ee9c

File tree

2 files changed

+22
-2
lines changed

2 files changed

+22
-2
lines changed

sdk/python/feast/infra/offline_stores/offline_store.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,26 @@ def metadata(self) -> Optional[RetrievalMetadata]:
163163
"""
164164
pass
165165

166+
def supports_remote_storage_export(self) -> bool:
167+
"""
168+
This method should return True if the RetrievalJob supports `to_remote_storage()`.
169+
"""
170+
return False
171+
172+
def to_remote_storage(self) -> List[str]:
173+
"""
174+
This method should export the result of this RetrievalJob to
175+
remote storage (such as S3, GCS, HDFS, etc).
176+
Implementations of this method should export the results as
177+
multiple parquet files, each file sized appropriately
178+
depending on how much data is being returned by the retrieval
179+
job.
180+
181+
Returns:
182+
A list of parquet file paths in remote storage.
183+
"""
184+
raise NotImplementedError()
185+
166186

167187
class OfflineStore(ABC):
168188
"""

sdk/python/tests/integration/offline_store/test_universal_historical_retrieval.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -280,7 +280,7 @@ def get_expected_training_df(
280280

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

@@ -410,7 +410,7 @@ def test_historical_features(environment, universal_data_sources, full_feature_n
410410

411411

412412
@pytest.mark.integration
413-
@pytest.mark.universal
413+
@pytest.mark.universal_offline_stores
414414
@pytest.mark.parametrize("full_feature_names", [True, False], ids=lambda v: str(v))
415415
def test_historical_features_with_shared_batch_source(
416416
environment, universal_data_sources, full_feature_names

0 commit comments

Comments
 (0)