Make overwriting optional

felixwang9817 · felixwang9817 · commit 8a1a6f0c649f · 2022-08-18T09:41:10.000-07:00
Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;
diff --git a/sdk/python/feast/feature_store.py b/sdk/python/feast/feature_store.py
@@ -1146,6 +1146,7 @@ def create_saved_dataset(
         storage: SavedDatasetStorage,
         tags: Optional[Dict[str, str]] = None,
         feature_service: Optional[FeatureService] = None,
+        allow_overwrite: bool = False,
     ) -> SavedDataset:
         """
         Execute provided retrieval job and persist its outcome in given storage.
@@ -1154,6 +1155,14 @@ def create_saved_dataset(
         Name for the saved dataset should be unique within project, since it's possible to overwrite previously stored dataset
         with the same name.
 
+        Args:
+            from_: The retrieval job whose result should be persisted.
+            name: The name of the saved dataset.
+            storage: The saved dataset storage object indicating where the result should be persisted.
+            tags (optional): A dictionary of key-value pairs to store arbitrary metadata.
+            feature_service (optional): The feature service that should be associated with this saved dataset.
+            allow_overwrite (optional): If True, the persisted result can overwrite an existing table or file.
+
         Returns:
             SavedDataset object with attached RetrievalJob
 
@@ -1186,7 +1195,7 @@ def create_saved_dataset(
         dataset.min_event_timestamp = from_.metadata.min_event_timestamp
         dataset.max_event_timestamp = from_.metadata.max_event_timestamp
 
-        from_.persist(storage)
+        from_.persist(storage=storage, allow_overwrite=allow_overwrite)
 
         dataset = dataset.with_retrieval_job(
             self._get_provider().retrieve_saved_dataset(
diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py
@@ -493,7 +493,7 @@ def _execute_query(
         block_until_done(client=self.client, bq_job=bq_job, timeout=timeout)
         return bq_job
 
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         assert isinstance(storage, SavedDatasetBigQueryStorage)
 
         self.to_bigquery(
diff --git a/sdk/python/feast/infra/offline_stores/contrib/athena_offline_store/athena.py b/sdk/python/feast/infra/offline_stores/contrib/athena_offline_store/athena.py
@@ -402,7 +402,7 @@ def _to_arrow_internal(self) -> pa.Table:
     def metadata(self) -> Optional[RetrievalMetadata]:
         return self._metadata
 
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         assert isinstance(storage, SavedDatasetAthenaStorage)
         self.to_athena(table_name=storage.athena_options.table)
 
diff --git a/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py b/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/postgres.py
@@ -297,7 +297,7 @@ def _to_arrow_internal(self) -> pa.Table:
     def metadata(self) -> Optional[RetrievalMetadata]:
         return self._metadata
 
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         assert isinstance(storage, SavedDatasetPostgreSQLStorage)
 
         df_to_postgres_table(
diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py
@@ -275,7 +275,7 @@ def _to_arrow_internal(self) -> pyarrow.Table:
             self.to_spark_df().write.parquet(temp_dir, mode="overwrite")
             return pq.read_table(temp_dir)
 
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         """
         Run the retrieval and persist the results in the same offline store used for read.
         Please note the persisting is done only within the scope of the spark session.
diff --git a/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino.py b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/trino.py
@@ -126,7 +126,7 @@ def to_trino(
         self._client.execute_query(query_text=query)
         return destination_table
 
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         """
         Run the retrieval and persist the results in the same offline store used for read.
         """
diff --git a/sdk/python/feast/infra/offline_stores/file.py b/sdk/python/feast/infra/offline_stores/file.py
@@ -88,11 +88,11 @@ def _to_arrow_internal(self):
         df = self.evaluation_function().compute()
         return pyarrow.Table.from_pandas(df)
 
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         assert isinstance(storage, SavedDatasetFileStorage)
 
         # Check if the specified location already exists.
-        if os.path.exists(storage.file_options.uri):
+        if not allow_overwrite and os.path.exists(storage.file_options.uri):
             raise SavedDatasetLocationAlreadyExists(location=storage.file_options.uri)
 
         filesystem, path = FileSource.create_filesystem_and_path(
diff --git a/sdk/python/feast/infra/offline_stores/offline_store.py b/sdk/python/feast/infra/offline_stores/offline_store.py
@@ -173,13 +173,15 @@ def on_demand_feature_views(self) -> List[OnDemandFeatureView]:
         pass
 
     @abstractmethod
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         """
         Synchronously executes the underlying query and persists the result in the same offline store
         at the specified destination.
 
-        Currently does not prevent overwriting a pre-existing location in the offline store, although
-        individual implementations may do so. Eventually all implementations should prevent overwriting.
+        Args:
+            storage: The saved dataset storage object specifying where the result should be persisted.
+            allow_overwrite: If True, a pre-existing location (e.g. table or file) can be overwritten.
+                Currently not all individual offline store implementations make use of this parameter.
         """
         pass
 
diff --git a/sdk/python/feast/infra/offline_stores/redshift.py b/sdk/python/feast/infra/offline_stores/redshift.py
@@ -483,7 +483,7 @@ def to_redshift(self, table_name: str) -> None:
                 query,
             )
 
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         assert isinstance(storage, SavedDatasetRedshiftStorage)
         self.to_redshift(table_name=storage.redshift_options.table)
 
diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py
@@ -460,7 +460,7 @@ def to_arrow_chunks(self, arrow_options: Optional[Dict] = None) -> Optional[List
 
         return arrow_batches
 
-    def persist(self, storage: SavedDatasetStorage):
+    def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         assert isinstance(storage, SavedDatasetSnowflakeStorage)
         self.to_snowflake(table_name=storage.snowflake_options.table)
 
diff --git a/sdk/python/tests/integration/e2e/test_validation.py b/sdk/python/tests/integration/e2e/test_validation.py
@@ -65,6 +65,7 @@ def test_historical_retrieval_with_validation(environment, universal_data_source
         from_=reference_job,
         name="my_training_dataset",
         storage=environment.data_source_creator.create_saved_dataset_destination(),
+        allow_overwrite=True,
     )
     saved_dataset = store.get_saved_dataset("my_training_dataset")
 
@@ -95,6 +96,7 @@ def test_historical_retrieval_fails_on_validation(environment, universal_data_so
         from_=reference_job,
         name="my_other_dataset",
         storage=environment.data_source_creator.create_saved_dataset_destination(),
+        allow_overwrite=True,
     )
 
     job = store.get_historical_features(
@@ -172,6 +174,7 @@ def test_logged_features_validation(environment, universal_data_sources):
         ),
         name="reference_for_validating_logged_features",
         storage=environment.data_source_creator.create_saved_dataset_destination(),
+        allow_overwrite=True,
     )
 
     log_source_df = store.get_historical_features(
@@ -245,6 +248,7 @@ def test_e2e_validation_via_cli(environment, universal_data_sources):
         from_=retrieval_job,
         name="reference_for_validating_logged_features",
         storage=environment.data_source_creator.create_saved_dataset_destination(),
+        allow_overwrite=True,
     )
     reference = saved_dataset.as_reference(
         name="test_reference", profiler=configurable_profiler
diff --git a/sdk/python/tests/integration/offline_store/test_universal_historical_retrieval.py b/sdk/python/tests/integration/offline_store/test_universal_historical_retrieval.py
@@ -381,6 +381,7 @@ def test_historical_features_persisting(
         name="saved_dataset",
         storage=environment.data_source_creator.create_saved_dataset_destination(),
         tags={"env": "test"},
+        allow_overwrite=True,
     )
 
     event_timestamp = DEFAULT_ENTITY_DF_EVENT_TIMESTAMP_COL

Original file line number	Diff line number	Diff line change
`@@ -483,7 +483,7 @@ def to_redshift(self, table_name: str) -> None:`
`483`	`483`	`query,`
`484`	`484`	`)`
`485`	`485`
`486`		`- def persist(self, storage: SavedDatasetStorage):`
	`486`	`+ def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):`
`487`	`487`	`assert isinstance(storage, SavedDatasetRedshiftStorage)`
`488`	`488`	`self.to_redshift(table_name=storage.redshift_options.table)`
`489`	`489`