fix: Using repo_config parameter in teardown to allow for feature-sto…

…re-yaml overrides (#4413) * fix: using repo_config parameter in teardown to allow for feature-store-yaml overrides Signed-off-by: Dan Baron <[email protected]> * fix: fixing linting and formatting issues in tests Signed-off-by: Dan Baron <[email protected]> * fix: removing unnecessary Path object construction Signed-off-by: Dan Baron <[email protected]> --------- Signed-off-by: Dan Baron <[email protected]>
feast-dev · Aug 16, 2024 · 0baeeb5 · 0baeeb5
1 parent 721ec74
commit 0baeeb5
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/sdk/python/feast/repo_operations.py b/sdk/python/feast/repo_operations.py
@@ -359,7 +359,7 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
 
 def teardown(repo_config: RepoConfig, repo_path: Optional[str]):
     # Cannot pass in both repo_path and repo_config to FeatureStore.
-    feature_store = FeatureStore(repo_path=repo_path, config=None)
+    feature_store = FeatureStore(repo_path=repo_path, config=repo_config)
     feature_store.teardown()
 
 

diff --git a/sdk/python/feast/transformation/pandas_transformation.py b/sdk/python/feast/transformation/pandas_transformation.py
@@ -1,5 +1,4 @@
-from types import FunctionType
-from typing import Any
+from typing import Any, Callable
 
 import dill
 import pandas as pd
@@ -15,7 +14,7 @@
 
 
 class PandasTransformation:
-    def __init__(self, udf: FunctionType, udf_string: str = ""):
+    def __init__(self, udf: Callable[[Any], Any], udf_string: str = ""):
         """
         Creates an PandasTransformation object.
 
@@ -30,11 +29,11 @@ def __init__(self, udf: FunctionType, udf_string: str = ""):
     def transform_arrow(
         self, pa_table: pyarrow.Table, features: list[Field]
     ) -> pyarrow.Table:
-        output_df_pandas = self.udf.__call__(pa_table.to_pandas())
+        output_df_pandas = self.udf(pa_table.to_pandas())
         return pyarrow.Table.from_pandas(output_df_pandas)
 
     def transform(self, input_df: pd.DataFrame) -> pd.DataFrame:
-        return self.udf.__call__(input_df)
+        return self.udf(input_df)
 
     def infer_features(self, random_input: dict[str, list[Any]]) -> list[Field]:
         df = pd.DataFrame.from_dict(random_input)

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
@@ -5,6 +5,7 @@
 import tempfile
 import uuid
 from pathlib import Path
+from subprocess import Popen
 from typing import Any, Dict, List, Optional
 
 import pandas as pd
@@ -367,7 +368,7 @@ class RemoteOfflineStoreDataSourceCreator(FileDataSourceCreator):
     def __init__(self, project_name: str, *args, **kwargs):
         super().__init__(project_name)
         self.server_port: int = 0
-        self.proc = None
+        self.proc: Optional[Popen[bytes]] = None
 
     def setup(self, registry: RegistryConfig):
         parent_offline_config = super().create_offline_store_config()
@@ -382,13 +383,13 @@ def setup(self, registry: RegistryConfig):
         repo_path = Path(tempfile.mkdtemp())
         with open(repo_path / "feature_store.yaml", "w") as outfile:
             yaml.dump(config.model_dump(by_alias=True), outfile)
-        repo_path = str(repo_path.resolve())
+        repo_path = repo_path.resolve()
 
         self.server_port = free_port()
         host = "0.0.0.0"
         cmd = [
             "feast",
-            "-c" + repo_path,
+            "-c" + str(repo_path),
             "serve_offline",
             "--host",
             host,