Add FirstK aggregation (#553)

Add FirstK aggregation in codebase along the lines of LastK aggregation.

.wordlist.txt

@@ -37,6 +37,7 @@ Eval
 eval
 endswith
 FennelDataAccessRole
+FirstK
 Flink
 Flink's
 GCP

docs/api.yml

@@ -76,6 +76,7 @@ sidebar:
  - "api-reference/aggregations/count"
  - "api-reference/aggregations/distinct"
  - "api-reference/aggregations/lastk"
+ - "api-reference/aggregations/firstk"
  - "api-reference/aggregations/max"
  - "api-reference/aggregations/min"
  - "api-reference/aggregations/stddev"
@@ -171,4 +172,4 @@ sidebar:
  - slug: "api-reference/expectations"
  title: "Expectations"
  pages:
- - "api-reference/expectations"
+ - "api-reference/expectations"

docs/examples/api-reference/aggregations/firstk.py

@@ -0,0 +1,150 @@
+from datetime import datetime
+from typing import List
+
+import pandas as pd
+import pytest
+
+from fennel.testing import mock
+
+__owner__ = "satwant@fennel.ai"
+
+
+@mock
+def test_basic(client):
+ # docsnip basic
+ from fennel.datasets import dataset, field, pipeline, Dataset, FirstK
+ from fennel.dtypes import Continuous
+ from fennel.lib import inputs
+ from fennel.connectors import source, Webhook
+
+ webhook = Webhook(name="webhook")
+
+ @source(webhook.endpoint("Transaction"), disorder="14d", cdc="append")
+ @dataset
+ class Transaction:
+ uid: int
+ amount: int
+ timestamp: datetime
+
+ @dataset(index=True)
+ class Aggregated:
+ uid: int = field(key=True)
+ amounts: List[int]
+ timestamp: datetime
+
+ @pipeline
+ @inputs(Transaction)
+ def firstk_pipeline(cls, ds: Dataset):
+ return ds.groupby("uid").aggregate(
+ # docsnip-highlight start
+ amounts=FirstK(
+ of="amount",
+ limit=10,
+ dedup=False,
+ window=Continuous("1d"),
+ ),
+ # docsnip-highlight end
+ )
+
+ # /docsnip
+ client.commit(message="msg", datasets=[Transaction, Aggregated])
+ # log some rows to the transaction dataset
+ client.log(
+ "webhook",
+ "Transaction",
+ pd.DataFrame(
+ [
+ {
+ "uid": 1,
+ "vendor": "A",
+ "amount": 10,
+ "timestamp": "2021-01-01T00:00:00",
+ },
+ {
+ "uid": 1,
+ "vendor": "B",
+ "amount": 20,
+ "timestamp": "2021-01-02T00:00:00",
+ },
+ {
+ "uid": 2,
+ "vendor": "A",
+ "amount": 30,
+ "timestamp": "2021-01-03T00:00:00",
+ },
+ {
+ "uid": 2,
+ "vendor": "B",
+ "amount": 40,
+ "timestamp": "2021-01-04T00:00:00",
+ },
+ {
+ "uid": 3,
+ "vendor": "A",
+ "amount": 50,
+ "timestamp": "2021-01-05T00:00:00",
+ },
+ {
+ "uid": 3,
+ "vendor": "B",
+ "amount": 60,
+ "timestamp": "2021-01-06T00:00:00",
+ },
+ ]
+ ),
+ )
+
+ # do lookup on the Aggregated dataset
+ ts = pd.Series(
+ [
+ datetime(2021, 1, 6, 0, 0, 0),
+ datetime(2021, 1, 6, 0, 0, 0),
+ datetime(2021, 1, 6, 0, 0, 0),
+ ]
+ )
+ df, found = Aggregated.lookup(ts, uid=pd.Series([1, 2, 3]))
+ assert found.tolist() == [True, True, True]
+ assert df["uid"].tolist() == [1, 2, 3]
+ assert df["amounts"].tolist() == [[], [], [50, 60]]
+
+
+@mock
+def test_invalid_type(client):
+ with pytest.raises(Exception):
+ # docsnip incorrect_type
+ from fennel.datasets import dataset, field, pipeline, Dataset, FirstK
+ from fennel.dtypes import Continuous
+ from fennel.lib import inputs
+ from fennel.connectors import source, Webhook
+
+ webhook = Webhook(name="webhook")
+
+ @source(webhook.endpoint("Transaction"), disorder="14d", cdc="append")
+ @dataset
+ class Transaction:
+ uid: int
+ amount: int
+ timestamp: datetime
+
+ @dataset
+ class Aggregated:
+ uid: int = field(key=True)
+ # docsnip-highlight next-line
+ amounts: int # should be List[int]
+ timestamp: datetime
+
+ @pipeline
+ @inputs(Transaction)
+ def bad_pipeline(cls, ds: Dataset):
+ return ds.groupby("uid").aggregate(
+ # docsnip-highlight start
+ amounts=FirstK(
+ of="amount",
+ limit=10,
+ dedup=False,
+ window=Continuous("1d"),
+ ),
+ # docsnip-highlight end
+ )
+
+ # /docsnip

docs/pages/api-reference/aggregations/firstk.md

@@ -0,0 +1,63 @@
+---
+title: FirstK
+order: 0
+status: published
+---
+### FirstK
+Aggregation to computes a rolling list of the earliest values for each group 
+within a window. 
+
+#### Parameters
+<Expandable title="of" type="str">
+Name of the field in the input dataset over which the aggregation should be 
+computed.
+</Expandable>
+
+<Expandable title="window" type="Window">
+The continuous window within which aggregation needs to be computed. Possible 
+values are `"forever"` or any [time duration](/api-reference/data-types/duration).
+</Expandable>
+
+<Expandable title="into_field" type="str">
+The name of the field in the output dataset that should store the result of this
+aggregation. This field is expected to be of type `List[T]` where `T` is the type
+of the field denoted by `of`.
+</Expandable>
+
+<Expandable title="limit" type="int">
+Since storing all the values for a group can get costly, FirstK expects a 
+`limit` to be specified which denotes the maximum size of the list that should 
+be maintained at any point.
+</Expandable>
+
+<Expandable title="dedup" type="bool">
+If set to True, only distinct values are stored else values stored in the first
+can have duplicates too.
+</Expandable>
+
+<pre snippet="api-reference/aggregations/firstk#basic" status="success" 
+ message="FirstK in window of 1 day">
+</pre>
+
+#### Returns
+<Expandable type="List[T]">
+Stores the result of the aggregation in the appropriate field of the output 
+dataset. 
+</Expandable>
+
+
+#### Errors
+<Expandable title="Incorrect output type">
+The column denoted by `into_field` in the output dataset must be of type `List[T]`
+where T is the type of the column denoted by `of` in the input dataset. Commit error
+is raised if this is not the case.
+</Expandable>
+
+:::warning
+Storing the full set of values and maintaining order between them can get costly, 
+so use this aggregation only when needed.
+:::
+
+<pre snippet="api-reference/aggregations/firstk#incorrect_type" status="error" 
+ message="amounts should be of type List[int], not int">
+</pre>

fennel/CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## [1.5.28] - 2024-09-21
+- Add FirstK aggregation
+
 ## [1.5.27] - 2024-09-20
 - Fix schema validation for int64 columns during mock client. Earlier, float64
  was considered valid for optional[int64] fields to account for typing limitation

fennel/__init__.py

@@ -9,6 +9,7 @@
  Distinct,
  Sum,
  LastK,
+ FirstK,
  Min,
  Max,
  Average,

fennel/client_tests/test_social_network.py

@@ -5,7 +5,7 @@
 import pytest
 
 import fennel._vendor.requests as requests
-from fennel import LastK
+from fennel import LastK, FirstK
 from fennel.connectors import source, Webhook
 from fennel.datasets import dataset, field, Dataset, pipeline, Count
 from fennel.dtypes import regex, oneof, Continuous
@@ -166,6 +166,27 @@ def last_viewed_post(cls, view_data: Dataset):
  )
 
 
+@meta(owner="ml-eng@myspace.com")
+@dataset(index=True)
+class FirstViewedPostByAgg:
+ user_id: str = field(key=True)
+ post_id: List[int]
+ time_stamp: datetime
+
+ @pipeline
+ @inputs(ViewData)
+ def first_viewed_post(cls, view_data: Dataset):
+ return view_data.groupby("user_id").aggregate(
+ FirstK(
+ into_field="post_id",
+ of="post_id",
+ window=Continuous("forever"),
+ limit=1,
+ dedup=False,
+ )
+ )
+
+
 # --- Featuresets ---
 
 
@@ -188,6 +209,9 @@ class UserFeatures:
  last_viewed_post2: List[int] = F(
  LastViewedPostByAgg.post_id, default=[-1] # type: ignore
  )
+ first_viewed_post: List[int] = F(
+ FirstViewedPostByAgg.post_id, default=[-1] # type: ignore
+ )
 
  @extractor(deps=[UserViewsDataset]) # type: ignore
  @inputs(Request.user_id)
@@ -234,6 +258,7 @@ def test_social_network(client):
  UserCategoryDataset,
  LastViewedPost,
  LastViewedPostByAgg,
+ FirstViewedPostByAgg,
  ],
  featuresets=[Request, UserFeatures],
  )
@@ -307,6 +332,11 @@ def test_social_network(client):
  assert last_post_viewed == [936609766, 735291550]
  assert last_post_viewed2 == last_post_viewed
 
+ first_post_viewed = [
+ x[0] for x in feature_df["UserFeatures.first_viewed_post"].to_list()
+ ]
+ assert first_post_viewed == [508698801, 43094523]
+
  if client.is_integration_client():
  return
  df = client.get_dataset_df("UserCategoryDataset")
@@ -427,6 +457,7 @@ def test_social_network_with_mock_log(client):
  UserCategoryDataset,
  LastViewedPost,
  LastViewedPostByAgg,
+ FirstViewedPostByAgg,
  ],
  featuresets=[Request, UserFeatures],
  )
@@ -494,5 +525,10 @@ def test_social_network_with_mock_log(client):
  assert last_post_viewed == [936609766, 735291550]
  assert last_post_viewed2 == last_post_viewed
 
+ first_post_viewed = [
+ x[0] for x in feature_df["UserFeatures.first_viewed_post"].to_list()
+ ]
+ assert first_post_viewed == [508698801, 43094523]
+
  df = client.get_dataset_df("UserCategoryDataset")
  assert df.shape == (1998, 4)