fix: test errors

MatsMoll · Nov 9, 2023 · 450a6a8 · 450a6a8
1 parent 4de5621
commit 450a6a8
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 21 deletions.
diff --git a/aligned/schemas/transformation.py b/aligned/schemas/transformation.py
@@ -56,8 +56,7 @@ def gracefull_transformation(
     transformation: Callable[[pd.Series], pd.Series],
 ) -> pd.Series:
     result = pd.Series(np.repeat(np.nan, repeats=is_valid_mask.shape[0]))
-    result.loc[is_valid_mask] = transformation(df.loc[is_valid_mask])
-    return result
+    return result.mask(is_valid_mask, transformation(df.loc[is_valid_mask]))
 
 
 class PsqlTransformation:

diff --git a/conftest.py b/conftest.py
@@ -419,12 +419,7 @@ class TitanicPassenger(FeatureView):
         sex = String().accepted_values(['male', 'female'])
         survived = Bool().description('If the passenger survived')
 
-        sibsp = (
-            Int32()
-            .lower_bound(0, is_inclusive=True)
-            .upper_bound(20, is_inclusive=True)
-            .description('Number of siblings on titanic')
-        )
+        sibsp = Int32().lower_bound(0).upper_bound(20).description('Number of siblings on titanic')
 
         cabin = String()
 
@@ -480,12 +475,7 @@ class TitanicPassenger(FeatureView):
         sex = String().accepted_values(['male', 'female'])
         survived = Bool().description('If the passenger survived')
 
-        sibsp = (
-            Int32()
-            .lower_bound(0, is_inclusive=True)
-            .upper_bound(20, is_inclusive=True)
-            .description('Number of siblings on titanic')
-        )
+        sibsp = Int32().lower_bound(0).upper_bound(20).description('Number of siblings on titanic')
 
         cabin = String()
 
@@ -616,12 +606,7 @@ class TitanicPassenger(FeatureView):
             embedding_size=50, storage=redis.index(name='name_embedding_index'), metadata=[age, sex]
         )
 
-        sibsp = (
-            Int32()
-            .lower_bound(0, is_inclusive=True)
-            .upper_bound(20, is_inclusive=True)
-            .description('Number of siblings on titanic')
-        )
+        sibsp = Int32().lower_bound(0).upper_bound(20).description('Number of siblings on titanic')
 
         double_sibsp = sibsp * 2
         square_sibsp = sibsp * sibsp