Prevent overwriting existing file for file offline store

felixwang9817 · felixwang9817 · commit 76bbdf1af285 · 2022-08-16T16:07:39.000-07:00
Signed-off-by: Felix Wang &lt;wangfelix98@gmail.com&gt;
diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py
@@ -204,6 +204,11 @@ def __init__(
         )
 
 
+class SavedDatasetLocationAlreadyExists(Exception):
+    def __init__(self, location: str):
+        super().__init__(f"Saved dataset location {location} already exists.")
+
+
 class FeastOfflineStoreInvalidName(Exception):
     def __init__(self, offline_store_class_name: str):
         super().__init__(
diff --git a/sdk/python/feast/infra/offline_stores/contrib/athena_offline_store/tests/data_source.py b/sdk/python/feast/infra/offline_stores/contrib/athena_offline_store/tests/data_source.py
@@ -94,7 +94,9 @@ def create_data_source(
             data_source=self.offline_store_config.data_source,
         )
 
-    def create_saved_dataset_destination(self) -> SavedDatasetAthenaStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetAthenaStorage:
         table = self.get_prefixed_table_name(
             f"persisted_ds_{str(uuid.uuid4()).replace('-', '_')}"
         )
diff --git a/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/tests/data_source.py b/sdk/python/feast/infra/offline_stores/contrib/postgres_offline_store/tests/data_source.py
@@ -119,7 +119,9 @@ def create_online_store(self) -> Dict[str, str]:
             "password": POSTGRES_PASSWORD,
         }
 
-    def create_saved_dataset_destination(self):
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ):
         # FIXME: ...
         return None
 
diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/tests/data_source.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/tests/data_source.py
@@ -1,5 +1,5 @@
 import uuid
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import pandas as pd
 from pyspark import SparkConf
@@ -96,7 +96,9 @@ def create_data_source(
             field_mapping=field_mapping or {"ts_1": "ts"},
         )
 
-    def create_saved_dataset_destination(self) -> SavedDatasetSparkStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetSparkStorage:
         table = f"persisted_{str(uuid.uuid4()).replace('-', '_')}"
         return SavedDatasetSparkStorage(
             table=table, query=None, path=None, file_format=None
diff --git a/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/tests/data_source.py b/sdk/python/feast/infra/offline_stores/contrib/trino_offline_store/tests/data_source.py
@@ -105,7 +105,9 @@ def create_data_source(
             field_mapping=field_mapping or {"ts_1": "ts"},
         )
 
-    def create_saved_dataset_destination(self) -> SavedDatasetTrinoStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetTrinoStorage:
         table = self.get_prefixed_table_name(
             f"persisted_ds_{str(uuid.uuid4()).replace('-', '_')}"
         )
diff --git a/sdk/python/feast/infra/offline_stores/file.py b/sdk/python/feast/infra/offline_stores/file.py
@@ -1,3 +1,4 @@
+import os
 import uuid
 from datetime import datetime
 from pathlib import Path
@@ -11,13 +12,16 @@
 import pytz
 from pydantic.typing import Literal
 
-from feast import FileSource, OnDemandFeatureView
 from feast.data_source import DataSource
-from feast.errors import FeastJoinKeysDuringMaterialization
+from feast.errors import (
+    FeastJoinKeysDuringMaterialization,
+    SavedDatasetLocationAlreadyExists,
+)
 from feast.feature_logging import LoggingConfig, LoggingSource
 from feast.feature_view import DUMMY_ENTITY_ID, DUMMY_ENTITY_VAL, FeatureView
 from feast.infra.offline_stores.file_source import (
     FileLoggingDestination,
+    FileSource,
     SavedDatasetFileStorage,
 )
 from feast.infra.offline_stores.offline_store import (
@@ -30,6 +34,7 @@
     get_pyarrow_schema_from_batch_source,
 )
 from feast.infra.registry.base_registry import BaseRegistry
+from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.repo_config import FeastConfigBaseModel, RepoConfig
 from feast.saved_dataset import SavedDatasetStorage
 from feast.usage import log_exceptions_and_usage
@@ -85,6 +90,11 @@ def _to_arrow_internal(self):
 
     def persist(self, storage: SavedDatasetStorage):
         assert isinstance(storage, SavedDatasetFileStorage)
+
+        # Check if the specified location already exists.
+        if os.path.exists(storage.file_options.uri):
+            raise SavedDatasetLocationAlreadyExists(location=storage.file_options.uri)
+
         filesystem, path = FileSource.create_filesystem_and_path(
             storage.file_options.uri,
             storage.file_options.s3_endpoint_override,
diff --git a/sdk/python/feast/infra/offline_stores/offline_store.py b/sdk/python/feast/infra/offline_stores/offline_store.py
@@ -174,7 +174,13 @@ def on_demand_feature_views(self) -> List[OnDemandFeatureView]:
 
     @abstractmethod
     def persist(self, storage: SavedDatasetStorage):
-        """Synchronously executes the underlying query and persists the result in the same offline store."""
+        """
+        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.
+        """
         pass
 
     @property
diff --git a/sdk/python/tests/integration/feature_repos/universal/data_source_creator.py b/sdk/python/tests/integration/feature_repos/universal/data_source_creator.py
@@ -49,7 +49,13 @@ def create_offline_store_config(self) -> FeastConfigBaseModel:
         ...
 
     @abstractmethod
-    def create_saved_dataset_destination(self) -> SavedDatasetStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetStorage:
+        """
+        Creates a saved dataset destination. If data_source is specified, uses the location of that
+        data source as the destination for the saved dataset.
+        """
         ...
 
     def create_logged_features_destination(self) -> LoggingDestination:
diff --git a/sdk/python/tests/integration/feature_repos/universal/data_sources/bigquery.py b/sdk/python/tests/integration/feature_repos/universal/data_sources/bigquery.py
@@ -95,7 +95,9 @@ def create_data_source(
             field_mapping=field_mapping or {"ts_1": "ts"},
         )
 
-    def create_saved_dataset_destination(self) -> SavedDatasetBigQueryStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetBigQueryStorage:
         table = self.get_prefixed_table_name(
             f"persisted_{str(uuid.uuid4()).replace('-', '_')}"
         )
diff --git a/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py b/sdk/python/tests/integration/feature_repos/universal/data_sources/file.py
@@ -59,7 +59,17 @@ def create_data_source(
             field_mapping=field_mapping or {"ts_1": "ts"},
         )
 
-    def create_saved_dataset_destination(self) -> SavedDatasetFileStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetFileStorage:
+        if data_source:
+            assert isinstance(data_source, FileSource)
+            return SavedDatasetFileStorage(
+                path=data_source.path,
+                file_format=ParquetFormat(),
+                s3_endpoint_override=None,
+            )
+
         d = tempfile.mkdtemp(prefix=self.project_name)
         self.dirs.append(d)
         return SavedDatasetFileStorage(
@@ -154,7 +164,9 @@ def create_data_source(
             s3_endpoint_override=f"http://{host}:{port}",
         )
 
-    def create_saved_dataset_destination(self) -> SavedDatasetFileStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetFileStorage:
         port = self.minio.get_exposed_port("9000")
         host = self.minio.get_container_host_ip()
 
diff --git a/sdk/python/tests/integration/feature_repos/universal/data_sources/redshift.py b/sdk/python/tests/integration/feature_repos/universal/data_sources/redshift.py
@@ -78,7 +78,9 @@ def create_data_source(
             database=self.offline_store_config.database,
         )
 
-    def create_saved_dataset_destination(self) -> SavedDatasetRedshiftStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetRedshiftStorage:
         table = self.get_prefixed_table_name(
             f"persisted_ds_{str(uuid.uuid4()).replace('-', '_')}"
         )
diff --git a/sdk/python/tests/integration/feature_repos/universal/data_sources/snowflake.py b/sdk/python/tests/integration/feature_repos/universal/data_sources/snowflake.py
@@ -66,7 +66,9 @@ def create_data_source(
             warehouse=self.offline_store_config.warehouse,
         )
 
-    def create_saved_dataset_destination(self) -> SavedDatasetSnowflakeStorage:
+    def create_saved_dataset_destination(
+        self, data_source: Optional[DataSource] = None
+    ) -> SavedDatasetSnowflakeStorage:
         table = self.get_prefixed_table_name(
             f"persisted_ds_{str(uuid.uuid4()).replace('-', '_')}"
         )
diff --git a/sdk/python/tests/integration/offline_store/test_persist.py b/sdk/python/tests/integration/offline_store/test_persist.py
@@ -0,0 +1,50 @@
+import pytest
+
+from feast.errors import SavedDatasetLocationAlreadyExists
+from tests.integration.feature_repos.repo_configuration import (
+    construct_universal_feature_views,
+)
+from tests.integration.feature_repos.universal.entities import (
+    customer,
+    driver,
+    location,
+)
+
+
+@pytest.mark.integration
+@pytest.mark.universal_offline_stores(only=["file"])
+def test_persist_does_not_overwrite(environment, universal_data_sources):
+    """
+    Tests that the persist method does not overwrite an existing location in the offline store.
+
+    This test currently is only run against the file offline store as it is the only implementation
+    that prevents overwriting. As more offline stores add this check, they should be added to this test.
+    """
+    store = environment.feature_store
+    entities, datasets, data_sources = universal_data_sources
+    feature_views = construct_universal_feature_views(data_sources)
+    store.apply([driver(), customer(), location(), *feature_views.values()])
+
+    features = [
+        "customer_profile:current_balance",
+        "customer_profile:avg_passenger_count",
+        "customer_profile:lifetime_trip_count",
+    ]
+
+    entity_df = datasets.entity_df.drop(
+        columns=["order_id", "origin_id", "destination_id"]
+    )
+    job = store.get_historical_features(
+        entity_df=entity_df,
+        features=features,
+    )
+
+    with pytest.raises(SavedDatasetLocationAlreadyExists):
+        # This should fail since persisting to a preexisting location is not allowed.
+        store.create_saved_dataset(
+            from_=job,
+            name="my_training_dataset",
+            storage=environment.data_source_creator.create_saved_dataset_destination(
+                data_source=data_sources.customer
+            ),
+        )

Original file line number	Diff line number	Diff line change
`@@ -94,7 +94,9 @@ def create_data_source(`
`94`	`94`	`data_source=self.offline_store_config.data_source,`
`95`	`95`	`)`
`96`	`96`
`97`		`- def create_saved_dataset_destination(self) -> SavedDatasetAthenaStorage:`
	`97`	`+ def create_saved_dataset_destination(`
	`98`	`+ self, data_source: Optional[DataSource] = None`
	`99`	`+ ) -> SavedDatasetAthenaStorage:`
`98`	`100`	`table = self.get_prefixed_table_name(`
`99`	`101`	`f"persisted_ds_{str(uuid.uuid4()).replace('-', '_')}"`
`100`	`102`	`)`
Original file line number	Diff line number	Diff line change
`@@ -119,7 +119,9 @@ def create_online_store(self) -> Dict[str, str]:`
`119`	`119`	`"password": POSTGRES_PASSWORD,`
`120`	`120`	`}`
`121`	`121`
`122`		`- def create_saved_dataset_destination(self):`
	`122`	`+ def create_saved_dataset_destination(`
	`123`	`+ self, data_source: Optional[DataSource] = None`
	`124`	`+ ):`
`123`	`125`	`# FIXME: ...`
`124`	`126`	`return None`
`125`	`127`
Original file line number	Diff line number	Diff line change
`@@ -105,7 +105,9 @@ def create_data_source(`
`105`	`105`	`field_mapping=field_mapping or {"ts_1": "ts"},`
`106`	`106`	`)`
`107`	`107`
`108`		`- def create_saved_dataset_destination(self) -> SavedDatasetTrinoStorage:`
	`108`	`+ def create_saved_dataset_destination(`
	`109`	`+ self, data_source: Optional[DataSource] = None`
	`110`	`+ ) -> SavedDatasetTrinoStorage:`
`109`	`111`	`table = self.get_prefixed_table_name(`
`110`	`112`	`f"persisted_ds_{str(uuid.uuid4()).replace('-', '_')}"`
`111`	`113`	`)`
Original file line number	Diff line number	Diff line change
`@@ -95,7 +95,9 @@ def create_data_source(`
`95`	`95`	`field_mapping=field_mapping or {"ts_1": "ts"},`
`96`	`96`	`)`
`97`	`97`
`98`		`- def create_saved_dataset_destination(self) -> SavedDatasetBigQueryStorage:`
	`98`	`+ def create_saved_dataset_destination(`
	`99`	`+ self, data_source: Optional[DataSource] = None`
	`100`	`+ ) -> SavedDatasetBigQueryStorage:`
`99`	`101`	`table = self.get_prefixed_table_name(`
`100`	`102`	`f"persisted_{str(uuid.uuid4()).replace('-', '_')}"`
`101`	`103`	`)`