[FIX] Add casting to pyarrow for using expression in assign. (#562)

fennel/CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## [1.5.22] - 2024-09-12
+- Add casting to pyarrow for using expression in assign.
+
 ## [1.5.21] - 2024-09-12
 - Raise an error if expectations are defined on terminal datasets.
 

fennel/client_tests/test_dataset.py

@@ -4794,3 +4794,60 @@ def test_exponential_aggregation(client):
         assert (
             abs(df["rating_exp_agg_7d"].iloc[i] - expected_result_agg[i]) < 1e-3
         ), f"{df['rating_exp_agg_7d'].iloc[i]} != {expected_result_agg[i]}"
+
+
+@pytest.mark.integration
+@mock
+def test_assign_with_aggregation(client):
+
+    @source(webhook.endpoint("A"), disorder="14d", cdc="append")
+    @dataset
+    class A:
+        user_id: int
+        value: int
+        ts: datetime
+
+    @dataset(index=True)
+    class B:
+        user_id: int = field(key=True)
+        const: str = field(key=True)
+        sum: int
+        count: int
+        ts: datetime
+
+        @pipeline
+        @inputs(A)
+        def pipeline(cls, event: Dataset):
+            return (
+                event.assign(const=lit("1").astype(str))
+                .groupby("user_id", "const")
+                .aggregate(
+                    sum=Sum(
+                        of="value",
+                        window=Continuous("forever"),
+                    ),
+                    count=Count(
+                        of="value", window=Continuous("forever"), unique=False
+                    ),
+                )
+            )
+
+    client.commit(datasets=[A, B], message="test")
+
+    now = datetime.now(timezone.utc)
+    df = pd.DataFrame(
+        {
+            "user_id": [1, 2, 3],
+            "value": [1, 2, 3],
+            "ts": [now, now, now],
+        }
+    )
+    client.log("fennel_webhook", "A", df)
+    client.sleep()
+
+    results, _ = client.lookup(
+        B,
+        keys=pd.DataFrame({"user_id": [1, 2, 3], "const": [1, 1, 1]}),
+    )
+    assert results["sum"].tolist() == [1, 2, 3]
+    assert results["count"].tolist() == [1, 1, 1]

fennel/testing/executor.py

@@ -835,9 +835,6 @@ def visitAssign(self, obj):
 
                 # Check the schema of the column
                 validate_field_in_df(field, df, self.cur_pipeline_name)
-
-                # Cast to arrow dtype
-                df = cast_df_to_arrow_dtype(df, fields)
             except Exception as e:
                 raise Exception(
                     f"Error in assign node for column `{obj.column}` for pipeline "
@@ -858,6 +855,8 @@ def visitAssign(self, obj):
                         f"Error in assign node for column `{col}` for pipeline "
                         f"`{self.cur_pipeline_name}`, {e}"
                     )
+        # Cast to arrow dtype
+        df = cast_df_to_arrow_dtype(df, fields)
         return NodeRet(
             df,
             input_ret.timestamp_field,