Cleanup and some new dataset features

MatsMoll · Dec 25, 2023 · 895df34 · 895df34
1 parent a0bbbce
commit 895df34
Show file tree

Hide file tree

Showing 27 changed files with 970 additions and 323 deletions.
diff --git a/aligned/active_learning/write_policy.py b/aligned/active_learning/write_policy.py
@@ -36,7 +36,7 @@ class ActiveLearningSampleSizePolicy(ActiveLearningWritePolicy):
 
     async def write(self, data: pl.LazyFrame, model: Model):
 
-        if not model.dataset_folder:
+        if not model.dataset_store:
             logger.info(
                 'Found no dataset folder. Therefore, no data will be written to an active learning dataset.'
             )
@@ -55,7 +55,7 @@ async def write(self, data: pl.LazyFrame, model: Model):
             dataset_subfolder = Path(self.dataset_folder_name) / str(self.write_timestamp)
             logger.info(f'Writing active learning data to {dataset_subfolder}')
 
-            dataset = model.dataset_folder.file_at(dataset_subfolder / self.dataset_file_name)
+            dataset = model.dataset_store.file_at(dataset_subfolder / self.dataset_file_name)
             await dataset.write(self.current_frame.write_csv().encode('utf-8'))
             self.unsaved_size = 0
 

diff --git a/aligned/compiler/feature_factory.py b/aligned/compiler/feature_factory.py
@@ -34,7 +34,6 @@
 from aligned.schemas.vector_storage import VectorStorage
 
 if TYPE_CHECKING:
-    from aligned.compiler.transformation_factory import FillNaStrategy
     from aligned.sources.s3 import AwsS3Config
 
 
@@ -426,19 +425,14 @@ def add_values(feature: FeatureFactory) -> None:
     def copy_type(self: T) -> T:
         raise NotImplementedError()
 
-    def fill_na(self: T, value: FillNaStrategy | Any) -> T:
-
-        from aligned.compiler.transformation_factory import (
-            ConstantFillNaStrategy,
-            FillMissingFactory,
-            FillNaStrategy,
-        )
+    def fill_na(self: T, value: FeatureFactory | Any) -> T:
+        from aligned.compiler.transformation_factory import FillMissingFactory
 
         instance: FeatureFactory = self.copy_type()  # type: ignore [attr-defined]
-        if isinstance(value, FillNaStrategy):
-            instance.transformation = FillMissingFactory(self, value)
-        else:
-            instance.transformation = FillMissingFactory(self, ConstantFillNaStrategy(value))
+        if not isinstance(value, FeatureFactory):
+            value = LiteralValue.from_value(value)
+
+        instance.transformation = FillMissingFactory(self, value)  # type: ignore [attr-defined]
         return instance  # type: ignore [return-value]
 
     def transformed_using_features_pandas(