diff --git a/fennel/client_tests/test_featureset.py b/fennel/client_tests/test_featureset.py index 1e043629..b88538d0 100644 --- a/fennel/client_tests/test_featureset.py +++ b/fennel/client_tests/test_featureset.py @@ -20,7 +20,7 @@ expectations, expect_column_values_to_be_between, ) -from fennel.expr import col, lit, when +from fennel.expr import col, lit, when, repeat from fennel.testing import mock, log ################################################################################ @@ -1326,11 +1326,19 @@ class UserInfoListLookup: country: List[Optional[str]] = F(UserInfoDataset.country) name: List[str] = F(UserInfoDataset.name, default="Unknown") + @featureset + class UserInfoListLookup2: + user_id_request: int = F(Request.user_id) + user_id: List[int] = F(repeat(col("user_id_request"), 5)) + country: List[Optional[str]] = F(UserInfoDataset.country) + name: List[str] = F(UserInfoDataset.name, default="Unknown") + client.commit( datasets=[UserInfoDataset], - featuresets=[UserInfoListLookup], + featuresets=[UserInfoListLookup, UserInfoListLookup2, Request], message="Initial commit", ) + client.sleep() now = datetime.now(timezone.utc) data = [ @@ -1370,6 +1378,20 @@ class UserInfoListLookup: ["Monica", "Rahul", "Unknown"], ] + feature_df2 = client.query( + outputs=[UserInfoListLookup2], + inputs=[Request.user_id], + input_dataframe=pd.DataFrame( + {"Request.user_id": [18232, 12345]} + ), + ) + + assert feature_df2.shape == (2, 4) + assert feature_df2["UserInfoListLookup2.country"].tolist() == [ + ["USA", "USA", "USA", "USA", "USA"], + [pd.NA, pd.NA, pd.NA, pd.NA, pd.NA], + ] + @pytest.mark.integration @mock @@ -1465,7 +1487,7 @@ class KeyedFeatureset: client.commit( datasets=[KeyedDataset], featuresets=[KeyedFeatureset], message="msg" ) - + client.sleep() # Log data now = datetime.now(timezone.utc) data = [ diff --git a/fennel/testing/mock_client.py b/fennel/testing/mock_client.py index e7bc94dd..17eedea3 100644 --- a/fennel/testing/mock_client.py +++ b/fennel/testing/mock_client.py @@ -527,6 +527,7 @@ def _transform_input_dataframe_from_inputs( input_dataframe[input_col], col_type ) except Exception as e: + print(input_dataframe) raise Exception( f"Error casting input dataframe column `{input_col}` for feature `{feature.fqn_}`, " f"dtype: `{feature.dtype}`: {e}"