Add fields parameter in join to fetch a subset of fields from rhs (#535)

* Add fields parameter in join to fetch a subset of fields from rhs

* Add right fields filtering in join to mock executor

* Remove a check disallowing right_fields from containing right's timestamp

* Add more checks and a test

* Update proto generated files and fix test.

* Add an integration test in social network scenario

* Address comments

* Rebased

* Fix when rhs fields has timestamp included

---------

Co-authored-by: Aditya Nambiar <aditya.nambiar007@gmail.com>

docs/pages/api-reference/operators/join.md

@@ -48,7 +48,8 @@ sides must have the same data types.
 <Expandable title="within" type="Tuple[Duration, Duration]" defaultVal='("forever", "0s")'>
 Optional kwarg specifying the time window relative to the left side timestamp 
 within which the join should be performed. This can be seen as adding another
-condition to join like `WHERE left_time - d1 < right_time AND right_time < left_time + d1`
+condition to join like `WHERE left_time - d1 < right_time AND right_time < left_time + d2`
+where (d1, d2) = within.
 - The first value in the tuple represents how far back in time should a join
  happen. The term "forever" means that we can go infinitely back in time 
  when searching for an event to join from the left-hand side data.

fennel/CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## [1.5.19] - 2024-09-06
+- Add ability to specify fields in join
+
 ## [1.5.18] - 2024-09-05
 - Struct initializer + arrow fixes + type promotion in assign
 

fennel/client_tests/test_dataset.py

@@ -1397,6 +1397,151 @@ def pipeline_join(cls, info: Dataset, sale: Dataset):
  ]
 
 
+@meta(owner="test@test.com")
+@source(
+ webhook.endpoint("MovieRevenueWithRightFields"),
+ disorder="14d",
+ cdc="upsert",
+)
+@dataset(index=True)
+class MovieRevenueWithRightFields:
+ movie: oneof(str, ["Jumanji", "Titanic", "RaOne", "ABC"]) = field( # type: ignore
+ key=True
+ )
+ revenue: int
+ extra_field: int
+ t: datetime
+
+
+@meta(owner="satwant@fennel.ai")
+@dataset(index=True)
+class MovieStatsWithRightFields:
+ movie: oneof(str, ["Jumanji", "Titanic", "RaOne", "ABC"]) = field( # type: ignore
+ key=True
+ )
+ rating: float
+ revenue_in_millions: float
+ t: Optional[datetime]
+ ts: datetime = field(timestamp=True)
+
+ @pipeline
+ @inputs(MovieRating, MovieRevenueWithRightFields)
+ def pipeline_join(cls, rating: Dataset, revenue: Dataset):
+ def to_millions(df: pd.DataFrame) -> pd.DataFrame:
+ df[str(cls.revenue_in_millions)] = df["revenue"] / 1000000
+ df[str(cls.revenue_in_millions)].fillna(-1, inplace=True)
+ return df[
+ [
+ str(cls.movie),
+ str(cls.t),
+ str(cls.ts),
+ str(cls.revenue_in_millions),
+ str(cls.rating),
+ ]
+ ]
+
+ rating = rating.rename({"t": "ts"}) # type: ignore
+ c = rating.join(
+ revenue, how="left", on=[str(cls.movie)], fields=["revenue", "t"]
+ )
+ # Transform provides additional columns which will be filtered out.
+ return c.transform(
+ to_millions,
+ schema={
+ str(cls.movie): oneof(
+ str, ["Jumanji", "Titanic", "RaOne", "ABC"]
+ ),
+ str(cls.rating): float,
+ str(cls.t): Optional[datetime],
+ str(cls.ts): datetime,
+ str(cls.revenue_in_millions): float,
+ },
+ )
+
+
+class TestBasicJoinWithRightFields(unittest.TestCase):
+ @pytest.mark.integration
+ @mock
+ def test_basic_join_with_fields(self, client):
+ # # Sync the dataset
+ client.commit(
+ message="msg",
+ datasets=[
+ MovieRating,
+ MovieRevenueWithRightFields,
+ MovieStatsWithRightFields,
+ RatingActivity,
+ ],
+ )
+ now = datetime.now(timezone.utc)
+ one_hour_ago = now - timedelta(hours=1)
+ data = [
+ ["Jumanji", 4, 343, 789, one_hour_ago],
+ ["Titanic", 5, 729, 1232, now],
+ ]
+ columns = ["movie", "rating", "num_ratings", "sum_ratings", "t"]
+ df = pd.DataFrame(data, columns=columns)
+ response = client.log("fennel_webhook", "MovieRating", df)
+ assert response.status_code == requests.codes.OK, response.json()
+
+ two_hours_ago = now - timedelta(hours=2)
+ data = [
+ ["Jumanji", 2000000, 1, two_hours_ago],
+ ["Titanic", 50000000, 2, now],
+ ]
+ columns = ["movie", "revenue", "extra_field", "t"]
+ df = pd.DataFrame(data, columns=columns)
+ response = client.log(
+ "fennel_webhook", "MovieRevenueWithRightFields", df
+ )
+ assert response.status_code == requests.codes.OK, response.json()
+ client.sleep()
+
+ # Do some lookups to verify pipeline_join is working as expected
+ keys = pd.DataFrame({"movie": ["Jumanji", "Titanic"]})
+ df, _ = client.lookup(
+ "MovieStatsWithRightFields",
+ keys=keys,
+ )
+ assert df.shape == (2, 5)
+ assert df["movie"].tolist() == ["Jumanji", "Titanic"]
+ assert df["rating"].tolist() == [4, 5]
+ assert df["revenue_in_millions"].tolist() == [2, 50]
+ assert df["t"].tolist() == [two_hours_ago, now]
+ assert "extra_field" not in df.columns
+
+ # Do some lookup at various timestamps in the past
+ ts = pd.Series([two_hours_ago, one_hour_ago, one_hour_ago, now])
+ keys = pd.DataFrame(
+ {"movie": ["Jumanji", "Jumanji", "Titanic", "Titanic"]}
+ )
+ df, _ = client.lookup(
+ "MovieStatsWithRightFields",
+ timestamps=ts,
+ keys=keys,
+ )
+ assert df.shape == (4, 5)
+ assert df["movie"].tolist() == [
+ "Jumanji",
+ "Jumanji",
+ "Titanic",
+ "Titanic",
+ ]
+ assert pd.isna(df["rating"].tolist()[0])
+ assert df["rating"].tolist()[1] == 4
+ assert pd.isna(df["rating"].tolist()[2])
+ assert df["rating"].tolist()[3] == 5
+ assert pd.isna(df["revenue_in_millions"].tolist()[0])
+ assert df["revenue_in_millions"].tolist()[1] == 2
+ assert pd.isna(df["revenue_in_millions"].tolist()[2])
+ assert df["revenue_in_millions"].tolist()[3] == 50
+ assert pd.isna(df["t"].tolist()[0])
+ assert df["t"].tolist()[1] == two_hours_ago
+ assert pd.isna(df["t"].tolist()[2])
+ assert df["t"].tolist()[3] == now
+ assert "extra_field" not in df.columns
+
+
 class TestBasicAggregate(unittest.TestCase):
  @pytest.mark.integration
  @mock

fennel/client_tests/test_social_network.py

@@ -41,6 +41,19 @@ class PostInfo:
  timestamp: datetime
 
 
+@source(
+ webhook.endpoint("PostInfoWithRightFields"), disorder="14d", cdc="upsert"
+)
+@dataset(index=True)
+@meta(owner="data-eng@myspace.com")
+class PostInfoWithRightFields:
+ title: str
+ category: str # type: ignore
+ post_id: int = field(key=True)
+ timestamp: datetime
+ extra_field: str
+
+
 @meta(owner="data-eng@myspace.com")
 @dataset
 @source(webhook.endpoint("ViewData"), disorder="14d", cdc="append")
@@ -100,6 +113,25 @@ def count_user_views(cls, view_data: Dataset, post_info: Dataset):
  )
 
 
+@meta(owner="ml-eng@myspace.com")
+@dataset(index=True)
+class UserCategoryDatasetWithRightFields:
+ user_id: str = field(key=True)
+ category: str = field(key=True)
+ num_views: int
+ time_stamp: datetime
+
+ @pipeline
+ @inputs(ViewData, PostInfoWithRightFields)
+ def count_user_views(cls, view_data: Dataset, post_info: Dataset):
+ post_info_enriched = view_data.join(
+ post_info, how="inner", on=["post_id"], fields=["title", "category"]
+ )
+ return post_info_enriched.groupby("user_id", "category").aggregate(
+ [Count(window=Continuous("6y 8s"), into_field="num_views")]
+ )
+
+
 @meta(owner="ml-eng@myspace.com")
 @dataset(index=True)
 class LastViewedPost:
@@ -166,6 +198,28 @@ def extract_user_views(cls, ts: pd.Series, user_ids: pd.Series):
  return views["num_views"]
 
 
+@meta(owner="feature-team@myspace.com")
+@featureset
+class UserFeaturesWithRightFields:
+ user_id: str = F(Request.user_id) # type: ignore
+ num_views: int
+ category: str = F(Request.category) # type: ignore
+ num_category_views: int = F(UserCategoryDatasetWithRightFields.num_views, default=0) # type: ignore
+ category_view_ratio: float = F(col("num_category_views") / col("num_views"))
+ last_viewed_post: int = F(LastViewedPost.post_id, default=-1) # type: ignore
+ last_viewed_post2: List[int] = F(
+ LastViewedPostByAgg.post_id, default=[-1] # type: ignore
+ )
+
+ @extractor(deps=[UserViewsDataset]) # type: ignore
+ @inputs(Request.user_id)
+ @outputs("num_views")
+ def extract_user_views(cls, ts: pd.Series, user_ids: pd.Series):
+ views, _ = UserViewsDataset.lookup(ts, user_id=user_ids) # type: ignore
+ views = views.fillna(0)
+ return views["num_views"]
+
+
 @pytest.mark.integration
 @mock
 def test_social_network(client):
@@ -259,6 +313,107 @@ def test_social_network(client):
  assert df.shape == (1998, 4)
 
 
+@pytest.mark.integration
+@mock
+def test_social_network_with_fields(client):
+ client.commit(
+ message="social network",
+ datasets=[
+ UserInfo,
+ PostInfoWithRightFields,
+ ViewData,
+ CityInfo,
+ UserViewsDataset,
+ UserCategoryDatasetWithRightFields,
+ LastViewedPost,
+ LastViewedPostByAgg,
+ ],
+ featuresets=[Request, UserFeaturesWithRightFields],
+ )
+ user_data_df = pd.read_csv("fennel/client_tests/data/user_data.csv")
+ post_data_df = pd.read_csv("fennel/client_tests/data/post_data.csv")
+ post_data_len = len(post_data_df.index)
+ post_data_df["extra_field"] = list(range(0, post_data_len))
+ view_data_df = pd.read_csv("fennel/client_tests/data/view_data_sampled.csv")
+ ts = "2018-01-01 00:00:00"
+ user_data_df["timestamp"] = ts
+ post_data_df["timestamp"] = ts
+ view_data_df["time_stamp"] = view_data_df["time_stamp"].apply(
+ lambda x: datetime.strptime(x, "%m/%d/%Y %H:%M %p")
+ )
+ # # Advance all timestamps by 6 years
+ user_data_df["timestamp"] = pd.to_datetime(
+ user_data_df["timestamp"]
+ ) + pd.DateOffset(years=4)
+ post_data_df["timestamp"] = pd.to_datetime(
+ post_data_df["timestamp"]
+ ) + pd.DateOffset(years=4)
+ view_data_df["time_stamp"] = view_data_df["time_stamp"] + pd.DateOffset(
+ years=4
+ )
+
+ res = client.log("fennel_webhook", "UserInfo", user_data_df)
+ assert res.status_code == requests.codes.OK, res.json()
+ res = client.log("fennel_webhook", "PostInfoWithRightFields", post_data_df)
+ assert res.status_code == requests.codes.OK, res.json()
+ res = client.log("fennel_webhook", "ViewData", view_data_df)
+ assert res.status_code == requests.codes.OK, res.json()
+
+ if client.is_integration_client():
+ client.sleep(120)
+
+ keys = pd.DataFrame(
+ {
+ "city": ["Wufeng", "Coyaima", "San Angelo"],
+ "gender": ["Male", "Male", "Female"],
+ }
+ )
+
+ df, found = client.lookup(
+ "CityInfo",
+ keys=keys,
+ )
+ assert found.to_list() == [True, True, True]
+
+ feature_df = client.query(
+ outputs=[UserFeaturesWithRightFields],
+ inputs=[Request.user_id, Request.category],
+ input_dataframe=pd.DataFrame(
+ {
+ "Request.user_id": [
+ "5eece14efc13ae6609000000",
+ "5eece14efc13ae660900003c",
+ ],
+ "Request.category": ["banking", "programming"],
+ }
+ ),
+ )
+ assert (
+ feature_df["UserFeaturesWithRightFields.num_views"].to_list(),
+ feature_df["UserFeaturesWithRightFields.num_category_views"].to_list(),
+ feature_df["UserFeaturesWithRightFields.category_view_ratio"].to_list(),
+ ) == ([2, 4], [0, 1], [0.0, 0.25])
+
+ # Assert that both the last_viewed_post and last_viewed_post2 features are extracted correctly
+ last_post_viewed = feature_df[
+ "UserFeaturesWithRightFields.last_viewed_post"
+ ].to_list()
+ last_post_viewed2 = [
+ x[0]
+ for x in feature_df[
+ "UserFeaturesWithRightFields.last_viewed_post2"
+ ].to_list()
+ ]
+ assert last_post_viewed == [936609766, 735291550]
+ assert last_post_viewed2 == last_post_viewed
+
+ if client.is_integration_client():
+ return
+ df = client.get_dataset_df("UserCategoryDatasetWithRightFields")
+ assert "extra_field" not in df.columns
+ assert df.shape == (1998, 4)
+
+
 @mock
 def test_social_network_with_mock_log(client):
  client.commit(